Gaussian prediction model training method, robot reinforcement learning method and electronic equipment

Abstract

The invention provides a Gaussian prediction model training method, a robot reinforcement learning method and electronic equipment, and relates to the technical field of artificial intelligence. The method comprises the following steps: acquiring a first sample image of a first preset view angle and a second sample image of a first target view angle; performing parameter prediction by adopting a pre-training Gaussian prediction model according to the first sample image to obtain a first prediction Gaussian parameter and a second prediction Gaussian parameter; generating a first prediction image according to the second prediction Gaussian parameter; according to the first training loss of the first prediction image and the second sample image and the second training loss of the first prediction Gaussian parameter and the second prediction Gaussian parameter, the pre-training Gaussian prediction model is trained to obtain a target Gaussian prediction model, and the target Gaussian prediction model is used for being deployed on the robot. And performing reinforcement learning by generating the scene Gaussian parameters of the working scene, and determining the action of the robot. According to the method, high-quality environment representation parameters can be generated, and the reinforcement learning effect is improved.

Description

Technical Field

[0001] The present application relates to the field of artificial intelligence technology, and more specifically, to a Gaussian prediction model training method, a robot reinforcement learning method, and an electronic device. Background Art

[0002] In reinforcement learning (RL), especially in vision-based RL tasks, excellent environment representation parameters are crucial to learning results, and their quality directly affects the completion of learning tasks.

[0003] Existing vision-based RL usually uses explicit or implicit methods to represent the environment, such as images, points, voxels, and neural radiation fields. However, these representation methods have many defects. They either cannot describe complex local geometric structures, or are difficult to generalize to unseen scenes. Some also require foreground masks, which makes it impossible to provide RL with excellent environment representation parameters, reducing the effect of reinforcement learning. Summary of the invention

[0004] The purpose of this application is to provide a Gaussian prediction model training method, a robot reinforcement learning method and an electronic device to address the deficiencies in the above-mentioned prior art, so as to generate high-quality environment representation parameters and improve the effect of reinforcement learning.

[0005] To achieve the above purpose, the technical solution adopted in the embodiment of the present application is as follows: In a first aspect, an embodiment of the present application provides a Gaussian prediction model training method, the method comprising: Acquire a first sample image of a training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle; Perform parameter prediction using a pre-trained Gaussian prediction model according to the first sample image to obtain a first predicted Gaussian parameter and a second predicted Gaussian parameter, wherein the second predicted Gaussian parameter is an attribute parameter after denoising the first predicted Gaussian parameter; Generate a first predicted image of the first target perspective according to the second predicted Gaussian parameters; Calculating a first training loss according to the first predicted image and the second sample image; Calculating a second training loss according to the first predicted Gaussian parameter and the second predicted Gaussian parameter; According to the first training loss and the second training loss, the pre-trained Gaussian prediction model is trained to obtain a target Gaussian prediction model. The target Gaussian prediction model is used to be deployed on a robot so that the robot generates scene Gaussian parameters of the working scene according to the target Gaussian prediction model, and reinforcement learning is performed according to the scene Gaussian parameters to determine the actions of the robot.

[0006] Optionally, the pre-trained Gaussian prediction model includes: a pre-trained parameter prediction model and an initial parameter optimization model; the pre-trained Gaussian prediction model is used to perform parameter prediction according to the first sample image to obtain a first predicted Gaussian parameter and a second predicted Gaussian parameter, including: Perform parameter prediction using the pre-trained parameter prediction model according to the first sample image to obtain the first predicted Gaussian parameters; The first predicted Gaussian parameters are optimized using the initial parameter optimization model to obtain the second predicted Gaussian parameters.

[0007] Optionally, performing parameter prediction using the pre-trained parameter prediction model according to the first sample image to obtain the first predicted Gaussian parameter includes: Using the depth prediction module in the pre-trained parameter prediction model to obtain a sample feature map of the first sample image, and performing depth estimation on the sample feature map to obtain a depth image; Using the attribute prediction module in the pre-trained parameter prediction model to perform parameter prediction on the first sample image, the sample feature map, and the fused feature map corresponding to the depth image to obtain predicted rotation parameters, predicted scaling parameters, and predicted transparency parameters; Determining predicted three-dimensional position parameters according to the depth image; Determining a predicted color parameter according to the color value of the first sample image; The first predicted Gaussian parameters include: the predicted three-dimensional position parameters, the predicted rotation parameters, the predicted scaling parameters, the predicted color parameters and the predicted transparency parameters.

[0008] Optionally, the using a depth prediction module in the pre-trained parameter prediction model to perform depth estimation on the sample feature map to obtain a depth image includes: If the first sample image is a single sample image, a single image depth predictor is used to perform depth prediction on a sample feature map of the single sample image to obtain the depth image; If the first sample image is two sample images, a disparity prediction network is used to perform disparity prediction on sample feature maps of the two sample images to obtain the depth image.

[0009] Optionally, acquiring a first sample image of the training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle includes: Determine the second sample image from the sample image set of the training three-dimensional scene; According to the first target viewing angle of the second sample image, two sample images of adjacent viewing angles closest to the first target viewing angle are determined from the sample image set as the first sample images.

[0010] Optionally, before performing parameter prediction using a pre-trained Gaussian prediction model according to the first sample image to obtain first predicted Gaussian parameters and second predicted Gaussian parameters, the method further includes: Acquire a third sample image of a second preset viewing angle and a fourth sample image of a second target viewing angle; Perform parameter prediction using an initial Gaussian prediction model according to the fourth sample image to obtain third predicted Gaussian parameters; Generate a second predicted image of the second target perspective according to the third predicted Gaussian parameters; Calculating a third training loss according to the second predicted image and the fourth sample image; The depth prediction module in the initial Gaussian prediction model is pre-trained according to the third training loss to obtain the pre-trained Gaussian prediction model.

[0011] Optionally, the training the pre-trained Gaussian prediction model according to the first training loss and the second training loss to obtain a target Gaussian prediction model includes: According to the first training loss and the second training loss, other modules in the pre-trained Gaussian prediction model except the depth prediction module are trained to obtain a target Gaussian prediction model.

[0012] Optionally, the initial parameter optimization model is a graph neural network model, and the optimizing the first predicted Gaussian parameters using the initial parameter optimization model to obtain the second predicted Gaussian parameters includes: Generate a first predicted Gaussian point cloud according to the first predicted Gaussian parameters; The first predicted Gaussian point cloud is optimized using the graph neural network model to obtain a second predicted Gaussian point cloud, where the second predicted Gaussian point cloud is a point cloud represented by the second predicted Gaussian parameters; The step of generating a first predicted image of the first target perspective according to the second predicted Gaussian parameter includes: Rendering is performed according to the second predicted Gaussian point cloud to generate the first predicted image.

[0013] In a second aspect, an embodiment of the present application further provides a robot reinforcement learning method based on a Gaussian prediction model, the method comprising: Acquire a scene image of a preset working scene; Using a pre-trained target Gaussian prediction model to perform parameter prediction on the scene image to obtain scene Gaussian parameters of the preset working scene, wherein the target Gaussian prediction model is pre-trained using the Gaussian prediction model training method as described in any one of the first aspects; A reinforcement learning network is used to predict the robot's actions according to the Gaussian parameters of the scene, so as to control the robot to operate in the preset working scene.

[0014] In a third aspect, an embodiment of the present application provides a Gaussian prediction model training device, the device comprising: A first image acquisition module, used to acquire a first sample image of a training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle; A first parameter prediction module, configured to perform parameter prediction based on the first sample image using a pre-trained Gaussian prediction model to obtain a first predicted Gaussian parameter and a second predicted Gaussian parameter, wherein the second predicted Gaussian parameter is an attribute parameter after denoising the first predicted Gaussian parameter; An image generation module, configured to generate a first predicted image of the first target viewing angle according to the second predicted Gaussian parameters; A loss calculation module, used for calculating a first training loss according to the first predicted image and the second sample image; The loss calculation module is further used to calculate a second training loss according to the first predicted Gaussian parameter and the second predicted Gaussian parameter; A model training module is used to train the pre-trained Gaussian prediction model according to the first training loss and the second training loss to obtain a target Gaussian prediction model. The target Gaussian prediction model is used to be deployed on a robot so that the robot generates scene Gaussian parameters of a working scene according to the target Gaussian prediction model, performs reinforcement learning according to the scene Gaussian parameters, and determines the actions of the robot.

[0015] Optionally, the pre-trained Gaussian prediction model includes: a pre-trained parameter prediction model and an initial parameter optimization model; the first parameter prediction module includes: A first parameter prediction unit, configured to perform parameter prediction according to the first sample image using the pre-trained parameter prediction model to obtain the first predicted Gaussian parameter; The second parameter prediction unit is used to optimize the first predicted Gaussian parameters using the initial parameter optimization model to obtain the second predicted Gaussian parameters.

[0016] Optionally, the first parameter prediction unit is specifically used to use the depth prediction module in the pre-trained parameter prediction model to obtain the sample feature map of the first sample image, and perform depth estimation on the sample feature map to obtain a depth image; use the attribute prediction module in the pre-trained parameter prediction model to perform parameter prediction on the first sample image, the sample feature map and the fusion feature map corresponding to the depth image to obtain predicted rotation parameters, predicted scaling parameters and predicted transparency parameters; determine the predicted three-dimensional position parameters according to the depth image; determine the predicted color parameters according to the color value of the first sample image; wherein the first predicted Gaussian parameters include: the predicted three-dimensional position parameters, the predicted rotation parameters, the predicted scaling parameters, the predicted color parameters and the predicted transparency parameters.

[0017] Optionally, the first parameter prediction unit is specifically used to, if the first sample image is a single sample image, use a single image depth predictor to perform depth prediction on the sample feature map of the single sample image to obtain the depth image; if the first sample image is two sample images, use a disparity prediction network to perform disparity prediction on the sample feature maps of the two sample images to obtain the depth image.

[0018] Optionally, the first image acquisition module is specifically used to determine the second sample image from the sample image set of the training three-dimensional scene; based on the first target perspective of the second sample image, two sample images of adjacent perspectives closest to the first target perspective are determined from the sample image set as the first sample images.

[0019] Optionally, the first image acquisition module is further used to acquire a third sample image of a second preset viewing angle and a fourth sample image of a second target viewing angle; The first parameter prediction module is further used to perform parameter prediction according to the fourth sample image using an initial Gaussian prediction model to obtain a third predicted Gaussian parameter; The image generation module is further used to generate a second predicted image of the second target viewing angle according to the third predicted Gaussian parameter; The loss calculation module is further used to calculate a third training loss according to the second predicted image and the fourth sample image; The model training module is also used to pre-train the depth prediction module in the initial Gaussian prediction model according to the third training loss to obtain the pre-trained Gaussian prediction model.

[0020] Optionally, the model training module is specifically used to train other modules in the pre-trained Gaussian prediction model except the depth prediction module according to the first training loss and the second training loss to obtain a target Gaussian prediction model.

[0021] Optionally, the initial parameter optimization model is a graph neural network model, and the second parameter prediction unit is specifically used to generate a first predicted Gaussian point cloud according to the first predicted Gaussian parameters; the first predicted Gaussian point cloud is optimized using the graph neural network model to obtain a second predicted Gaussian point cloud, and the second predicted Gaussian point cloud is a point cloud represented by the second predicted Gaussian parameters; The image generation module is specifically used to perform rendering according to the second predicted Gaussian point cloud to generate the first predicted image.

[0022] In a fourth aspect, an embodiment of the present application further provides a robot reinforcement learning device based on a Gaussian prediction model, the device comprising: A second image acquisition module, used to acquire a scene image of a preset working scene; A second parameter prediction module, used to use a pre-trained target Gaussian prediction model to perform parameter prediction on the scene image to obtain scene Gaussian parameters of the preset working scene, wherein the target Gaussian prediction model is pre-trained using the Gaussian prediction model training method as described in any one of the first aspects; The control module is used to predict the robot's actions using a reinforcement learning network according to the Gaussian parameters of the scene, so as to control the robot to operate in the preset working scene.

[0023] In the fifth aspect, an embodiment of the present application also provides an electronic device, comprising: a processor, a storage medium and a bus, wherein the storage medium stores program instructions executable by the processor, and when the electronic device is running, the processor communicates with the storage medium through the bus, and the processor executes the program instructions to perform the steps of the Gaussian prediction model training method as described in any one of the first aspect, or the steps of the robot reinforcement learning method based on the Gaussian prediction model as described in the second aspect.

[0024] In a sixth aspect, an embodiment of the present application further provides a computer-readable storage medium, on which a computer program is stored. When the computer program is executed by a processor, the Gaussian prediction model training method as described in any one of the first aspects, or the robot reinforcement learning method based on the Gaussian prediction model as described in the second aspect, is executed.

[0025] The beneficial effects of this application are: The Gaussian prediction model training method, robot reinforcement learning method and electronic device provided in the present application train a Gaussian prediction model through sample images of a three-dimensional scene, so that the scene Gaussian parameters of the scene can be generated through the Gaussian prediction model, and the three-dimensional local geometric structure of the scene is described in detail through the scene Gaussian parameters, which facilitates the robot to perform reinforcement learning based on the scene Gaussian parameters, improves the robot's reinforcement learning effect, and ensures the robot's operation accuracy; and based on the Gaussian prediction model, Gaussian parameter prediction can be provided for robots in different scenes, which is generalizable and has a wide range of application scenarios. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] In order to more clearly illustrate the technical solutions of the embodiments of the present application, the drawings required for use in the embodiments will be briefly introduced below. It should be understood that the following drawings only show certain embodiments of the present application and therefore should not be regarded as limiting the scope. For ordinary technicians in this field, other related drawings can be obtained based on these drawings without paying creative work.

[0027] Figure 1 Schematic diagram of the process of the Gaussian prediction model training method provided in the embodiment of the present application Figure 1 ; Figure 2 A model architecture diagram provided for an embodiment of the present application; Figure 3 Schematic diagram of the process of the Gaussian prediction model training method provided in the embodiment of the present application Figure 2 ; Figure 4 Schematic diagram of the process of the Gaussian prediction model training method provided in the embodiment of the present application Figure 3 ; Figure 5 A schematic diagram of a process flow of a robot reinforcement learning method based on a Gaussian prediction model provided in an embodiment of the present application; Figure 6 A schematic diagram of the structure of a Gaussian prediction model training device provided in an embodiment of the present application; Figure 7 A schematic diagram of the structure of a robot reinforcement learning device based on a Gaussian prediction model provided in an embodiment of the present application; Figure 8 A schematic diagram of an electronic device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0028] In order to make the purpose, technical solutions and advantages of the embodiments of the present application clearer, the technical solutions in the embodiments of the present application will be clearly and completely described below in combination with the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, not all of the embodiments.

[0029] Therefore, the following detailed description of the embodiments of the present application provided in the accompanying drawings is not intended to limit the scope of the present application for which protection is sought, but merely represents selected embodiments of the present application. Based on the embodiments in the present application, all other embodiments obtained by ordinary technicians in the field without creative work are within the scope of protection of the present application.

[0030] In addition, the terms "first", "second", etc. in the specification and claims of the present application and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequential order. It should be understood that the data used in this way can be interchangeable where appropriate, so that the embodiments of the present application described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any of their variations are intended to cover non-exclusive inclusions, for example, a process, method, system, product or device comprising a series of steps or units is not necessarily limited to those steps or units clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.

[0031] It should be noted that, in the absence of conflict, the features in the embodiments of the present application may be combined with each other.

[0032] Figure 1 Schematic diagram of the process of the Gaussian prediction model training method provided in the embodiment of the present application Figure 1 ,like Figure 1 As shown, the method may include: S101, obtaining a first sample image of a training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle.

[0033] In this embodiment, when the robot is working in a three-dimensional scene, it needs to perform reinforcement learning based on the environmental representation parameters of the three-dimensional scene to determine the working actions. In order to provide the robot with accurate environmental representation parameters that can describe the three-dimensional scene, a Gaussian prediction model is planned to be trained.

[0034] Image acquisition is performed on a training three-dimensional scene to obtain sample images of the training three-dimensional scene at multiple perspectives to form a sample image set, and a first sample image of a first preset perspective and a second sample image of a first target perspective are obtained from the sample image set, wherein the first target perspective is the perspective of a predicted image to be generated, the second sample image of the first target perspective is used to supervise the training of a Gaussian prediction model, the first preset perspective is an adjacent perspective of the first target perspective, and the first sample image of the first preset perspective is used as a training sample of the Gaussian prediction model.

[0035] S102 . Perform parameter prediction using a pre-trained Gaussian prediction model according to the first sample image to obtain first predicted Gaussian parameters and second predicted Gaussian parameters, where the second predicted Gaussian parameters are attribute parameters after denoising the first predicted Gaussian parameters.

[0036] In this embodiment, the pre-trained Gaussian prediction model is a Gaussian prediction model that has been partially trained in advance. Some model parameters of the pre-trained Gaussian prediction model are frozen in this training, that is, some model parameters of the pre-trained Gaussian prediction model do not participate in the training during this training process. In this way, pre-training is performed before training, which can improve the training efficiency of the model.

[0037] Part of the model parameters common to different three-dimensional scenes in the Gaussian prediction model may be pre-trained.

[0038] The first sample image is input into a pre-trained Gaussian prediction model. The pre-trained Gaussian prediction model generates first predicted Gaussian parameters by performing feature extraction, encoding, decoding and other processes on the first sample image. The first predicted Gaussian parameters are used to represent Gaussian attribute parameters of the training three-dimensional scene.

[0039] After generating the first predicted Gaussian parameters, since this scheme needs to generate a first predicted image of the first target perspective based on the first sample image of the first preset perspective, in order to reduce the high-frequency noise caused by the differences between adjacent perspectives, the pre-trained Gaussian prediction model also optimizes the first predicted Gaussian parameters, filters the noise in the first predicted Gaussian parameters, and obtains smoothed second predicted Gaussian parameters.

[0040] S103 . Generate a first predicted image of a first target viewing angle according to the second predicted Gaussian parameters.

[0041] In this embodiment, the second predicted Gaussian parameters are environment representation parameters of the training three-dimensional scene, and rasterization rendering is performed based on the second predicted Gaussian parameters to obtain a first predicted image of the first target perspective.

[0042] Specifically, a first predicted image of the first target perspective may be generated according to the second predicted Gaussian parameters and the camera intrinsic parameters and the posture of the second sample image of the first target perspective.

[0043] S104: Calculate a first training loss according to the first predicted image and the second sample image.

[0044] In this embodiment, the first training loss is calculated according to the difference between each pixel point between the first predicted image and the second sample image.

[0045] For example, the first training loss may be calculated based on the square of the Euclidean distance between the first predicted image and the second sample image.

[0046] S105. Calculate a second training loss according to the first predicted Gaussian parameter and the second predicted Gaussian parameter.

[0047] In this embodiment, in order to supervise the optimization process of the pre-trained Gaussian prediction model generating the second predicted Gaussian parameters according to the first predicted Gaussian parameters, it is also necessary to calculate the second training loss between the first predicted Gaussian parameters and the second predicted Gaussian parameters.

[0048] For example, the second training loss can be calculated based on the square of the Euclidean distance between the first predicted Gaussian parameter and the second predicted Gaussian parameter.

[0049] S106. According to the first training loss and the second training loss, the pre-trained Gaussian prediction model is trained to obtain a target Gaussian prediction model. The target Gaussian prediction model is used to be deployed on the robot so that the humanoid robot generates scene Gaussian parameters of the working scene according to the target Gaussian prediction model, and reinforcement learning is performed according to the scene Gaussian parameters to determine the actions of the humanoid robot.

[0050] In this embodiment, the first training loss and the second training loss are weighted, and the model parameters other than the frozen parameters in the pre-trained prediction model are optimized according to the weighted training loss. The above steps are used to perform multiple rounds of training, and the training is stopped when the training rounds reach the preset rounds or the weighted training loss converges to obtain the target Gaussian prediction model.

[0051] After the training is completed, the target Gaussian prediction model is deployed in the robot. By acquiring images of the robot's working scene from any perspective, the target Gaussian prediction model is used to generate scene Gaussian parameters of the robot's working scene, i.e., environment representation parameters, so that the reinforcement learning network in the robot performs reinforcement learning based on the scene Gaussian parameters to determine the robot's next action.

[0052] The Gaussian prediction model training method provided in the above embodiment trains the Gaussian prediction model through sample images of the three-dimensional scene, so that the scene Gaussian parameters of the scene can be generated through the Gaussian prediction model, and the three-dimensional local geometric structure of the scene is described in detail through the scene Gaussian parameters, which is convenient for the robot to perform reinforcement learning based on the scene Gaussian parameters, improve the reinforcement learning effect of the robot, and ensure the operation accuracy of the robot; and based on the Gaussian prediction model, Gaussian parameter prediction can be provided for robots in different scenes, which is generalizable and has a wide range of application scenarios.

[0053] In one possible implementation, Figure 2 The model architecture diagram provided for the embodiment of the present application is as follows: Figure 2As shown, the pre-trained Gaussian prediction model may include: a pre-trained parameter prediction model and an initial parameter optimization model. The above S102 uses the pre-trained Gaussian prediction model to perform parameter prediction according to the first sample image to obtain the first predicted Gaussian parameter and the second predicted Gaussian parameter. The process may include: According to the first sample image, a pre-trained parameter prediction model is used to perform parameter prediction to obtain first predicted Gaussian parameters; and according to the first predicted Gaussian parameters, an initial parameter optimization model is used to perform optimization to obtain second predicted Gaussian parameters.

[0054] In this embodiment, the pre-trained Gaussian prediction model is divided into two parts: a pre-trained parameter prediction model and an initial parameter optimization model, wherein some model parameters in the pre-trained parameter prediction model are frozen after pre-training.

[0055] The first sample image is input into the pre-trained parameter prediction model, and the pre-trained parameter prediction model generates a first predicted Gaussian parameter by performing feature extraction, encoding, decoding and other processes on the first sample image.

[0056] The first preset Gaussian parameter is input into the initial parameter optimization model, and the initial parameter optimization model encodes and decodes the first predicted Gaussian parameter to reduce the high-frequency noise caused by the difference between adjacent views in the first predicted Gaussian parameter, and obtains the smoothed second predicted Gaussian parameter.

[0057] The Gaussian prediction model training method provided in the above embodiment optimizes the predicted Gaussian parameters output by the parameter prediction model based on the parameter optimization model, reduces the high-frequency noise caused by the differences between adjacent views, and improves the smoothness of the predicted Gaussian parameters, so that the robot's reinforcement learning network can control the robot's movements based on precise scene Gaussian parameters, thereby improving the robot's operation accuracy.

[0058] In one possible implementation, Figure 3 Schematic diagram of the process of the Gaussian prediction model training method provided in the embodiment of the present application Figure 2 ,like Figure 3 As shown, the process of performing parameter prediction using the pre-trained parameter prediction model according to the first sample image to obtain the first predicted Gaussian parameter may include: S201. Use a depth prediction module in a pre-trained parameter prediction model to obtain a sample feature map of a first sample image, and perform depth estimation on the sample feature map to obtain a depth image.

[0059] S202, using the attribute prediction module in the pre-trained parameter prediction model to perform parameter prediction on the first sample image, the sample feature map, and the fused feature map corresponding to the depth image to obtain predicted rotation parameters, predicted scaling parameters, and predicted transparency parameters.

[0060] S203: Determine predicted three-dimensional position parameters according to the depth image.

[0061] S204: Determine a predicted color parameter according to the color value of the first sample image.

[0062] The first predicted Gaussian parameters include: predicted three-dimensional position parameters, predicted rotation parameters, predicted scaling parameters, predicted color parameters and predicted transparency parameters.

[0063] In this embodiment, if Figure 2 As shown, the pre-trained parameter prediction model may include: a depth prediction module, a feature fusion module and an attribute prediction module, wherein the depth prediction model includes an encoder and a decoder, the encoder is a feature extractor, and the feature extractor is used to extract features of the first sample image to obtain a sample feature map of the first sample image.

[0064] The sample feature map of the first sample image is input into the decoder of the depth prediction module for depth estimation, and the depth value of each pixel in the first sample image is determined to obtain a depth image of the first sample image.

[0065] The first sample image, the sample feature map of the first sample image and the depth image are connected through a feature fusion module to obtain a fused feature map, wherein the feature fusion module can use an encoder to generate a fused feature map with full image resolution.

[0066] The fused feature map is input into the attribute prediction module, which includes multiple parameter decoders, namely a rotation parameter decoder, a scaling parameter decoder and a transparency parameter decoder. The fused feature map is decoded by the multiple parameter decoders to obtain predicted rotation parameters, predicted scaling parameters and predicted transparency parameters.

[0067] In some embodiments, multiple parameter decoders are each composed of a full convolutional layer with a 1*1 convolution kernel.

[0068] In addition to predicted rotation parameters, predicted scaling parameters and predicted transparency parameters, the first predicted Gaussian parameters also include predicted three-dimensional position parameters and predicted color parameters, wherein the predicted three-dimensional position parameters can be determined based on the depth image and the camera posture of the first sample image taken from the first target perspective, and the predicted color parameters can be directly determined based on the RGB color values ​​of each pixel in the first sample image.

[0069] The predicted rotation parameters and predicted scaling parameters represent the rotation information and scaling information of the Gaussian points corresponding to the object structure in the training three-dimensional scene. The predicted transparency parameters are used to represent the occlusion of light by the Gaussian points corresponding to the object structure in the training three-dimensional scene. The predicted three-dimensional position parameters are used to represent the position information of the Gaussian points corresponding to the object structure in the training three-dimensional scene. The predicted color parameters are used to represent the color information of the Gaussian points corresponding to the object structure in the training three-dimensional scene.

[0070] It should be noted that the above-mentioned prediction parameters are all pixel-by-pixel prediction parameters, and the predicted parameters are actually parameter maps.

[0071] In a possible implementation, the process of performing depth estimation on the sample feature map by using the depth prediction module in the pre-trained parameter prediction model in S202 to obtain the depth image may include: If the first sample image is a single sample image, a single image depth predictor is used to perform depth prediction on the sample feature map of the single sample image to obtain a depth image; if the first sample image is two sample images, a disparity prediction network is used to perform disparity prediction on the sample feature maps of the two sample images to obtain a depth image.

[0072] In this embodiment, a single sample image can be used to predict Gaussian parameters, or a multi-view image can be used to predict Gaussian parameters, wherein if the first sample image is a single sample image, the depth prediction module is a pre-trained single-image depth predictor, and the single-image depth predictor is used to perform depth prediction based on the sample feature map of the single sample image to obtain a depth image. If the first sample image is two sample images, the depth prediction module is a pre-trained disparity prediction network, and two identical feature extractors are used to extract features from the two sample images to obtain two sample feature maps, and the disparity prediction network is used to predict the disparity of each pixel in the two sample feature maps, and the depth value of each pixel is determined in combination with the maximum disparity value in the disparity data set of each pixel to obtain a depth image.

[0073] In some embodiments, if the first sample image is two sample images, the process of obtaining the first sample image of the training three-dimensional scene at the first preset perspective and the second sample image at the first target perspective in S101 may include: A second sample image is determined from a sample image set of a training three-dimensional scene; and according to a first target perspective of the second sample image, two sample images of adjacent perspectives closest to the first target perspective are determined from the sample image set as first sample images.

[0074] In this embodiment, the sample image set for training three-dimensional scenes includes sample images of multiple perspectives. First, a second sample image of a first target perspective is determined from the sample image set. Then, based on the distance between the first target perspective and the perspectives of other sample images, two sample images of two adjacent perspectives closest to the first target perspective are determined as first sample images.

[0075] Furthermore, if the first sample image is two sample images, the process of determining the predicted color parameter according to the color value of the first sample image may be determining the predicted color parameter according to the average value of the color values ​​of the two sample images.

[0076] In one possible implementation, Figure 4 Schematic diagram of the process of the Gaussian prediction model training method provided in the embodiment of the present application Figure 3 ,like Figure 4 As shown, before the above S102 uses the pre-trained Gaussian prediction model to perform parameter prediction according to the first sample image to obtain the first predicted Gaussian parameter and the second predicted Gaussian parameter, the method may further include: S301, acquiring a third sample image of a second preset viewing angle and a fourth sample image of a second target viewing angle.

[0077] S302 . Perform parameter prediction using an initial Gaussian prediction model according to the fourth sample image to obtain third predicted Gaussian parameters.

[0078] S303: Generate a second predicted image of a second target perspective according to the third predicted Gaussian parameters.

[0079] S304: Calculate a third training loss according to the second predicted image and the fourth sample image.

[0080] S305. Pre-train the depth prediction module in the initial Gaussian prediction model according to the third training loss to obtain a pre-trained Gaussian prediction model.

[0081] In this embodiment, the depth prediction module in the initial Gaussian prediction model needs to be pre-trained first, and the pre-training step is similar to the training step of S101-S106, except that the training loss in the pre-training step is only the third training loss between the fourth sample image and the second predicted image, and does not include the training loss between the two predicted Gaussian parameters. The third predicted Gaussian parameter is an optimized Gaussian parameter.

[0082] In some embodiments, an L1 loss function may be used to calculate the third training loss of the second predicted image and the fourth sample image.

[0083] The Gaussian prediction model training method provided in the above embodiment can improve the training efficiency of the model by pre-training the depth prediction module in the Gaussian prediction model.

[0084] In a possible implementation, the initial parameter optimization model is a graph neural network model, and the process of optimizing the first predicted Gaussian parameter using the initial parameter optimization model to obtain the second predicted Gaussian parameter may include: A first predicted Gaussian point cloud is generated according to the first predicted Gaussian parameters. The first predicted Gaussian point cloud is optimized by using a graph neural network model to obtain a second predicted Gaussian point cloud, and the second predicted Gaussian point cloud is a point cloud represented by the second predicted Gaussian parameters.

[0085] The process of generating the first predicted image of the first target perspective according to the second predicted Gaussian parameters in S103 may include: Rendering is performed according to the second predicted Gaussian point cloud to generate a first predicted image.

[0086] In this embodiment, the initial position of the point cloud is determined according to the predicted three-dimensional position parameters in the first predicted Gaussian parameters, the orientation of the initial position of the point cloud is adjusted according to the predicted rotation parameters, the scale of the point cloud is adjusted according to the predicted scaling parameters to obtain the position of the point cloud, and the color and opacity of the power source are defined according to the predicted color parameters and the predicted transparency parameters to obtain the first predicted Gaussian point cloud.

[0087] The first predicted Gaussian point cloud is input into the graph neural network model for optimization, the positions of each point cloud in the first predicted Gaussian point cloud are smoothed, and the high-frequency noise caused by the differences between different views is removed to obtain the second predicted Gaussian point cloud. Rasterization rendering is performed based on the second predicted Gaussian point cloud to generate the first predicted image.

[0088] In some embodiments, the graph neural network model may include an encoder and a decoder, the encoder may include multiple multi-layer perceptrons for outputting a multi-dimensional node feature map, and the number of layers of the decoder is the same as that of the encoder.

[0089] The Gaussian prediction model training method provided in the above embodiment adopts a graph neural network model to optimize the first predicted Gaussian parameters, and utilizes the smoothing characteristics of the graph neural network model to effectively reduce the high-frequency noise caused by view differences, thereby improving the accuracy of the second predicted Gaussian parameters.

[0090] In the Gaussian prediction model training method provided in the above embodiment, the embodiment of the present application also provides a robot reinforcement learning method based on the Gaussian prediction model. Figure 5 A schematic diagram of a robot reinforcement learning method based on a Gaussian prediction model provided in an embodiment of the present application is shown in FIG. Figure 5 As shown, the method may include: S401: Acquire a scene image of a preset working scene.

[0091] S402: Use a pre-trained target Gaussian prediction model to perform parameter prediction on the scene image to obtain scene Gaussian parameters of a preset working scene.

[0092] S403: using a reinforcement learning network to predict the robot's actions according to the Gaussian parameters of the scene, so as to control the robot to operate in a preset working scene.

[0093] In this embodiment, the target Gaussian prediction model is deployed on the robot. The robot is also provided with a camera, through which the scene image of the robot's working scene can be obtained. The target Gaussian prediction model can generate scene Gaussian parameters of the working scene based on the scene image, and generate a scene Gaussian point cloud based on the scene Gaussian parameters. The scene Gaussian point cloud is used to predict the robot's motion through a reinforcement learning network deployed on the robot, so as to control the robot to perform operations in the working scene.

[0094] The robot reinforcement learning method based on the Gaussian prediction model provided in the above embodiment adopts the Gaussian prediction model trained by Yuxuan to generate scene Gaussian parameters, so that the reinforcement learning network controls the robot to perform actions based on the scene Gaussian parameters. The high-quality scene Gaussian parameters improve the accuracy of the reinforcement learning network in predicting actions.

[0095] Based on the above method embodiments, an embodiment of the present application provides a Gaussian prediction model training device. Figure 6 A schematic diagram of the structure of the Gaussian prediction model training device provided in the embodiment of the present application is shown in FIG. Figure 6 As shown, the device may include: The first image acquisition module 501 is used to acquire a first sample image of a training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle; A first parameter prediction module 502 is used to perform parameter prediction based on the first sample image using a pre-trained Gaussian prediction model to obtain a first predicted Gaussian parameter and a second predicted Gaussian parameter, where the second predicted Gaussian parameter is an attribute parameter after denoising the first predicted Gaussian parameter; An image generation module 503 is used to generate a first predicted image of a first target viewing angle according to a second predicted Gaussian parameter; A loss calculation module 504, configured to calculate a first training loss according to the first predicted image and the second sample image; The loss calculation module 504 is further used to calculate the second training loss according to the first predicted Gaussian parameter and the second predicted Gaussian parameter; The model training module 505 is used to train the pre-trained Gaussian prediction model according to the first training loss and the second training loss to obtain a target Gaussian prediction model. The target Gaussian prediction model is used to be deployed on the robot so that the robot generates scene Gaussian parameters of the working scene according to the target Gaussian prediction model, performs reinforcement learning according to the scene Gaussian parameters, and determines the robot's actions.

[0096] Optionally, the pre-trained Gaussian prediction model includes: a pre-trained parameter prediction model and an initial parameter optimization model; the first parameter prediction module 502 includes: A first parameter prediction unit, configured to perform parameter prediction according to the first sample image using a pre-trained parameter prediction model to obtain a first predicted Gaussian parameter; The second parameter prediction unit is used to optimize the first predicted Gaussian parameters using an initial parameter optimization model to obtain second predicted Gaussian parameters.

[0097] Optionally, the first parameter prediction unit is specifically used to use the depth prediction module in the pre-trained parameter prediction model to obtain the sample feature map of the first sample image, and perform depth estimation on the sample feature map to obtain a depth image; use the attribute prediction module in the pre-trained parameter prediction model to perform parameter prediction on the first sample image, the sample feature map and the fusion feature map corresponding to the depth image to obtain predicted rotation parameters, predicted scaling parameters and predicted transparency parameters; determine the predicted three-dimensional position parameters according to the depth image; determine the predicted color parameters according to the color value of the first sample image; wherein the first predicted Gaussian parameters include: predicted three-dimensional position parameters, predicted rotation parameters, predicted scaling parameters, predicted color parameters and predicted transparency parameters.

[0098] Optionally, the first parameter prediction unit is specifically used to, if the first sample image is a single sample image, use a single image depth predictor to perform depth prediction on the sample feature map of the single sample image to obtain a depth image; if the first sample image is two sample images, use a disparity prediction network to perform disparity prediction on the sample feature maps of the two sample images to obtain a depth image.

[0099] Optionally, the first image acquisition module 501 is specifically used to determine a second sample image from a sample image set of a training three-dimensional scene; based on a first target perspective of the second sample image, two sample images of adjacent perspectives closest to the first target perspective are determined from the sample image set as first sample images.

[0100] Optionally, the first image acquisition module 501 is further configured to acquire a third sample image of a second preset viewing angle and a fourth sample image of a second target viewing angle; The first parameter prediction module 502 is further used to perform parameter prediction according to the fourth sample image using the initial Gaussian prediction model to obtain third predicted Gaussian parameters; The image generation module 503 is further used to generate a second predicted image of a second target viewing angle according to the third predicted Gaussian parameter; The loss calculation module 504 is further used to calculate a third training loss according to the second predicted image and the fourth sample image; The model training module 505 is also used to pre-train the depth prediction module in the initial Gaussian prediction model according to the third training loss to obtain a pre-trained Gaussian prediction model.

[0101] Optionally, the model training module 505 is specifically used to train other modules in the pre-trained Gaussian prediction model except the depth prediction module according to the first training loss and the second training loss to obtain a target Gaussian prediction model.

[0102] Optionally, the initial parameter optimization model is a graph neural network model, and the second parameter prediction unit is specifically used to generate a first predicted Gaussian point cloud according to the first predicted Gaussian parameters; the first predicted Gaussian point cloud is optimized using the graph neural network model to obtain a second predicted Gaussian point cloud, and the second predicted Gaussian point cloud is a point cloud represented by the second predicted Gaussian parameters; The image generation module 503 is specifically configured to perform rendering according to the second predicted Gaussian point cloud to generate a first predicted image.

[0103] The embodiment of the present application also provides a robot reinforcement learning device based on a Gaussian prediction model. Figure 7 A schematic diagram of the structure of a robot reinforcement learning device based on a Gaussian prediction model provided in an embodiment of the present application is shown in FIG. Figure 7 As shown, the device may include: The second image acquisition module 601 is used to acquire a scene image of a preset working scene; A second parameter prediction module 602 is used to use a pre-trained target Gaussian prediction model to perform parameter prediction on the scene image to obtain scene Gaussian parameters of a preset working scene, wherein the target Gaussian prediction model is pre-trained using any Gaussian prediction model training method as described in the first aspect; The control module 603 is used to predict the robot's actions using a reinforcement learning network according to the scene Gaussian parameters, so as to control the robot to operate in a preset working scene.

[0104] The above-mentioned device is used to execute the method provided by the aforementioned embodiment, and its implementation principle and technical effect are similar, which will not be repeated here.

[0105] The above modules may be one or more integrated circuits configured to implement the above methods, such as one or more application specific integrated circuits (ASICs), or one or more microprocessors, or one or more field programmable gate arrays (FPGAs). For another example, when a module is implemented in the form of a processing element scheduling program code, the processing element may be a general-purpose processor, such as a central processing unit (CPU) or other processor that can call program code. For another example, these modules may be integrated together and implemented in the form of a system-on-a-chip (SOC).

[0106] Figure 8 A schematic diagram of an electronic device provided in an embodiment of the present application, the electronic device 700 may include: a processor 701, a storage medium 702 and a bus, the storage medium 702 stores program instructions executable by the processor 701, when the electronic device 700 is running, the processor 701 communicates with the storage medium 702 via the bus, and the processor 701 executes the program instructions to execute the above method embodiment. The specific implementation method and technical effect are similar and will not be repeated here.

[0107] Optionally, the present application further provides a computer-readable storage medium, on which a computer program is stored, and when the computer program is executed by a processor, the above method embodiment is executed.

[0108] In the several embodiments provided in the present application, it should be understood that the disclosed devices and methods can be implemented in other ways. For example, the device embodiments described above are only schematic. For example, the division of units is only a logical function division. There may be other division methods in actual implementation, such as multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of devices or units, which can be electrical, mechanical or other forms.

[0109] The units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed on multiple network units. Some or all of the units may be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0110] In addition, each functional unit in each embodiment of the present application may be integrated into one processing unit, or each unit may exist physically separately, or two or more units may be integrated into one unit. The above-mentioned integrated unit may be implemented in the form of hardware or in the form of hardware plus software functional units.

[0111] The above-mentioned integrated unit implemented in the form of a software functional unit can be stored in a computer-readable storage medium. The above-mentioned software functional unit is stored in a storage medium, including a number of instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) or a processor (English: processor) to perform some steps of the methods of each embodiment of the present application. The aforementioned storage medium includes: U disk, mobile hard disk, read-only memory (English: Read-Only Memory, abbreviated: ROM), random access memory (English: Random Access Memory, abbreviated: RAM), disk or optical disk and other media that can store program codes.

[0112] The above is only a specific implementation of the present application, but the protection scope of the present application is not limited thereto. Any technician familiar with the technical field can easily think of changes or substitutions within the technical scope disclosed in the present application, which should be included in the protection scope of the present application. Therefore, the protection scope of the present application should be based on the protection scope of the claims.

Claims

1. A Gaussian prediction model training method, characterized in that: The method comprises: Acquire a first sample image of a training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle; Perform parameter prediction using a pre-trained Gaussian prediction model according to the first sample image to obtain a first predicted Gaussian parameter and a second predicted Gaussian parameter, wherein the second predicted Gaussian parameter is an attribute parameter after denoising the first predicted Gaussian parameter; Generate a first predicted image of the first target perspective according to the second predicted Gaussian parameters; Calculating a first training loss according to the first predicted image and the second sample image; Calculating a second training loss according to the first predicted Gaussian parameter and the second predicted Gaussian parameter; According to the first training loss and the second training loss, the pre-trained Gaussian prediction model is trained to obtain a target Gaussian prediction model. The target Gaussian prediction model is used to be deployed on a robot so that the robot generates scene Gaussian parameters of the working scene according to the target Gaussian prediction model, and reinforcement learning is performed according to the scene Gaussian parameters to determine the actions of the robot.

2. The method according to claim 1, characterized in that The pre-trained Gaussian prediction model includes: a pre-trained parameter prediction model and an initial parameter optimization model; the pre-trained Gaussian prediction model is used to perform parameter prediction according to the first sample image to obtain a first predicted Gaussian parameter and a second predicted Gaussian parameter, including: Perform parameter prediction using the pre-trained parameter prediction model according to the first sample image to obtain the first predicted Gaussian parameters; The first predicted Gaussian parameters are optimized using the initial parameter optimization model to obtain the second predicted Gaussian parameters.

3. The method according to claim 2, characterized in that The step of performing parameter prediction using the pre-trained parameter prediction model according to the first sample image to obtain the first predicted Gaussian parameter includes: Using the depth prediction module in the pre-trained parameter prediction model to obtain a sample feature map of the first sample image, and performing depth estimation on the sample feature map to obtain a depth image; Using the attribute prediction module in the pre-trained parameter prediction model to perform parameter prediction on the first sample image, the sample feature map, and the fused feature map corresponding to the depth image to obtain predicted rotation parameters, predicted scaling parameters, and predicted transparency parameters; Determining predicted three-dimensional position parameters according to the depth image; Determining a predicted color parameter according to the color value of the first sample image; The first predicted Gaussian parameters include: the predicted three-dimensional position parameters, the predicted rotation parameters, the predicted scaling parameters, the predicted color parameters and the predicted transparency parameters.

4. The method according to claim 3, characterized in that The step of using the depth prediction module in the pre-trained parameter prediction model to perform depth estimation on the sample feature map to obtain a depth image includes: If the first sample image is a single sample image, a single image depth predictor is used to perform depth prediction on a sample feature map of the single sample image to obtain the depth image; If the first sample image is two sample images, a disparity prediction network is used to perform disparity prediction on sample feature maps of the two sample images to obtain the depth image.

5. The method according to claim 4, characterized in that The step of acquiring a first sample image of a training three-dimensional scene at a first preset viewing angle and a second sample image at a first target viewing angle comprises: Determine the second sample image from the sample image set of the training three-dimensional scene; According to the first target viewing angle of the second sample image, two sample images of adjacent viewing angles closest to the first target viewing angle are determined from the sample image set as the first sample images.

6. The method according to claim 3, characterized in that Before performing parameter prediction using a pre-trained Gaussian prediction model according to the first sample image to obtain first predicted Gaussian parameters and second predicted Gaussian parameters, the method further includes: Acquire a third sample image of a second preset viewing angle and a fourth sample image of a second target viewing angle; Perform parameter prediction using an initial Gaussian prediction model according to the fourth sample image to obtain third predicted Gaussian parameters; Generate a second predicted image of the second target perspective according to the third predicted Gaussian parameters; Calculating a third training loss according to the second predicted image and the fourth sample image; The depth prediction module in the initial Gaussian prediction model is pre-trained according to the third training loss to obtain the pre-trained Gaussian prediction model.

7. The method according to claim 6, characterized in that The step of training the pre-trained Gaussian prediction model according to the first training loss and the second training loss to obtain a target Gaussian prediction model includes: According to the first training loss and the second training loss, other modules in the pre-trained Gaussian prediction model except the depth prediction module are trained to obtain a target Gaussian prediction model.

8. The method according to claim 2, characterized in that The initial parameter optimization model is a graph neural network model, and the first predicted Gaussian parameter is optimized using the initial parameter optimization model to obtain the second predicted Gaussian parameter, including: Generate a first predicted Gaussian point cloud according to the first predicted Gaussian parameters; The first predicted Gaussian point cloud is optimized by using the graph neural network model to obtain a second predicted Gaussian point cloud, where the second predicted Gaussian point cloud is a point cloud represented by the second predicted Gaussian parameters; The step of generating a first predicted image of the first target perspective according to the second predicted Gaussian parameter includes: Rendering is performed according to the second predicted Gaussian point cloud to generate the first predicted image.

9. A robot reinforcement learning method based on a Gaussian prediction model, characterized in that: The method comprises: Acquire a scene image of a preset working scene; Using a pre-trained target Gaussian prediction model to perform parameter prediction on the scene image to obtain scene Gaussian parameters of the preset working scene, wherein the target Gaussian prediction model is pre-trained using the Gaussian prediction model training method according to any one of claims 1 to 8; A reinforcement learning network is used to predict the robot's actions according to the Gaussian parameters of the scene, so as to control the robot to operate in the preset working scene.

10. An electronic device, characterized in that: include: A processor, a storage medium and a bus, wherein the storage medium stores program instructions executable by the processor, and when the electronic device is running, the processor communicates with the storage medium via the bus, and the processor executes the program instructions to execute the Gaussian prediction model training method as described in any one of claims 1 to 8, or the robot reinforcement learning method based on the Gaussian prediction model as described in claim 9.

Images