Reinforcement learning based joint 6d pose estimation method

Abstract

The application discloses a joint object 6D pose estimation method based on reinforcement learning, which is characterized in that the 6D pose estimation problem is converted into a partial point cloud registration task, a point cloud registration track is used as a sample to train a strategy network, the strategy network is used for aligning a source point cloud, namely a complete object in a standard space, to a target point cloud, namely a partial object observed in a camera space, then a state for training an intelligent agent is generated by using the trained strategy network, an action space and a reward function of reinforcement learning are constructed to train the intelligent agent, and finally, the sub-component pose is calculated by constructing and optimizing a joint object kinematics energy function. The application significantly improves the joint object pose estimation performance.

Description

Technical Field

[0001] This invention relates to a technology in the field of artificial intelligence, specifically to a 6D pose estimation method for joints based on reinforcement learning. Background Technology

[0002] Accurate 6D pose perception of jointed objects is crucial for many downstream tasks, such as robot manipulation, VR, and AR. Existing 6D pose estimation methods for jointed objects typically construct a deep neural network to densely predict the corresponding points in standard space from the observed point cloud; however, these dense prediction methods require excessive computation. Furthermore, the observed jointed object may encounter self-occlusion problems, which limits performance when estimating small parts. Summary of the Invention

[0003] To address the aforementioned shortcomings of existing technologies, this invention proposes a 6D pose estimation method for joints based on reinforcement learning. This method models the task as a reinforcement learning problem and then uses reinforcement learning to learn the partial point cloud registration trajectory, and iteratively estimates the 6D position of each component.

[0004] This invention is achieved through the following technical solution:

[0005] This invention relates to a 6D pose estimation method for jointed objects based on reinforcement learning. The 6D pose estimation problem is transformed into a partial point cloud registration task. The point cloud registration trajectory is used as a sample to train the policy network, which is used to align the source point cloud, i.e., the complete object in the standard space, to the target point cloud, i.e., the partial object observed in the camera space. Then, the trained policy network is used to generate the state for training the agent, constructing the action space and reward function of reinforcement learning to train the agent. Finally, the pose of the sub-parts is calculated by constructing and optimizing the kinematic energy function of the jointed object.

[0006] The policy initialization mentioned above refers to: constructing a PointNet++ structure as a policy network for imitation learning, processing the source object point cloud P and the target partial point cloud P', and outputting a serialized pose fitting trajectory Tr={(r1,t1),(r2,t2),…,(r k ,t k )}, where: R i and t i These represent the 3D rotation and translation at step i, respectively.

[0007] The policy network described above uniformly samples the true sequential fitting trajectory Tr'={(r'1,t'1),(r'2,t'2),…,(r' k ,t' kAfter that, the cross-entropy loss function is used to construct the training loss function L to train the policy network, specifically: Where: L r The rotational error, L t The translation error is λ. r and λ t The scaling factors were set to 0.3 and 0.7 respectively, because the translation fitting strategy is more difficult to converge.

[0008] The reinforcement learning algorithm comprises three main elements: state S, action space A, and reward function R. State S is represented by the feature vector output from the last layer of the PointNet++ architecture. Given state S, 12 actions are designed to construct the action space A. Six actions represent rotation operations: positive X-axis rotation, negative X-axis rotation, positive Y-axis rotation, negative Y-axis rotation, positive Z-axis rotation, and negative Z-axis rotation. The other six actions represent translation operations: positive X-axis translation, negative X-axis translation, positive Y-axis translation, negative Y-axis translation, positive Z-axis translation, and negative Z-axis translation. The reward function R is based on the point cloud P that the source object point cloud is currently transformed into at step i. i The Chamfer distance between the target point cloud P' and the target point cloud is designed as follows: Where: η is set to 0.5, S i S i-1 S and S' represent the shapes of the source point cloud at step i and step (i-1), and the corresponding states of the target point cloud, respectively. Using this method, a reinforcement learning agent can be trained to estimate the pose of the root components of an articulated object.

[0009] The kinematic energy function refers to the following: When estimating the pose of a sub-component of a joint object, the pose of each sub-component is first initialized with the pose {r0,t0} of the root component. Based on the information of each joint defined in the point cloud of the source object, and based on the kinematic structure of the input joint object, an energy function for the state of each joint is constructed. The energy function is optimized to find the optimal state of the joint. The pose of each sub-component is obtained by integrating the states of each joint.

[0010] This invention relates to a system for implementing the above method, comprising: a point cloud feature extraction unit, a reinforcement learning policy initialization unit, a reinforcement agent training unit, and a sub-component pose optimization unit, wherein: the point cloud feature extraction unit uses the PointNet++ structure to process the point cloud data at each training stage to obtain feature vectors that are convenient for subsequent processing; the reinforcement learning policy initialization unit uses imitation learning to initialize the policy network parameters based on the real sequential fitting trajectory; the reinforcement agent training unit uses a pre-defined state, action space, and reward function to deeply optimize the policy network and train an agent that can be used to generate point cloud registration trajectories; the sub-component pose optimization unit optimizes a pre-defined energy function to calculate the most likely pose of each sub-component of the object.

[0011] Technical effect

[0012] This invention transforms the 6D pose estimation problem into a partial point cloud registration task. It uses reinforcement learning to learn a point cloud registration trajectory to describe the process of aligning the source point cloud to the target point cloud. The PointNet++ architecture is used to process the point cloud data and extract feature vectors. Imitation learning is used to initialize the policy network, and trajectories sampled from real data are used as expert demonstrations for supervised learning. During the reinforcement training phase, rotation and translation operations performed on the point cloud in various directions are used as the action space, and a reward function is set based on the change in the Chamfer distance between the point cloud and the target point cloud at each step. Based on the kinematic structure of the joint object, an energy function for the joint state is constructed to optimize the pose of sub-components. Compared with existing technologies, this invention significantly improves the pose estimation performance of joint objects. This method achieves state-of-the-art performance on the PartNet-Mobility dataset. For some categories (such as ovens and washing machines), it achieves rotation estimation errors of 1.6°, 2.2°, 0.9°, and 1.1° with only 0.9 seconds of inference time. It also performs well on the ArtImage dataset, with rotation errors of only 2.8° and 3.1° in the laptop category, outperforming the baseline method OMAD by 2.6° and 1.2°. In sub-part pose estimation, this method achieves 3.6° and 3.5° on two sub-parts in the glasses category, significantly outperforming OMAD, indicating that it can address the self-occlusion problem in joint pose estimation tasks to some extent. Attached Figure Description

[0013] Figure 1 This is a flowchart of the present invention;

[0014] Figure 2 This is a schematic diagram of the system of the present invention;

[0015] Figure 3This is a schematic diagram of policy initialization based on imitation learning.

[0016] Figure 4 A schematic diagram of the action space during the reinforcement learning phase. Detailed Implementation

[0017] like Figure 1 As shown, this embodiment illustrates a 6D pose estimation method for joints based on reinforcement learning, comprising:

[0018] Step 1: Use the PointNet++ structure as a feature extractor to process the input point cloud structure during training and inference, that is, convert the point cloud into a 2048-dimensional feature vector for subsequent processing.

[0019] Step 2: Initialize the policy network using imitation learning, specifically including:

[0020] Step 2.1: Given the source object point cloud P and the target part point cloud P', use the feature extractor from Step 1 to convert both into feature vectors, and then use a multilayer fully connected network (MLP) to further convert the two feature vectors into a sequential pose fitting trajectory Tr={(r1,t1),(r2,t2),…,(r k ,t k )}, where: R i and t i These represent the 3D rotation and translation at step i, respectively. The structure formed by PointNet++ and MLP can be viewed as a shared policy network in the process of imitation learning and reinforcement learning.

[0021] Step 2.2: Using the true 6D pose of the object in the training data, take the object's pose in standard space (resting state) and its actual pose in real space as the start and end points, respectively, and generate a true sequential fitting trajectory Tr'={(r'1,t'1),(r'2,t'2),…,(r' k ,t' k )}.

[0022] Step 2.3: Use the cross-entropy loss function to construct a training loss function L to train the policy network: Where: L r The rotational error, L t The translation error is measured, CE represents the cross-entropy loss function, and λ r and λ t The scaling factors were set to 0.3 and 0.7 respectively, because the translation fitting strategy is more difficult to converge.

[0023] Step 3: Train the reinforcement learning agent to perform pose estimation, specifically including:

[0024] Step 3.1: The feature vector obtained by transforming the input point cloud through the feature extractor described in Step 1 is used as the state S in the reinforcement learning process.

[0025] Step 3.2: Construct the action space A for reinforcement learning, which contains 12 actions, representing positive rotation along the X-axis, negative rotation along the X-axis, positive rotation along the Y-axis, negative rotation along the Y-axis, positive rotation along the Z-axis, negative rotation along the Z-axis, positive translation along the X-axis, negative translation along the X-axis, positive translation along the Y-axis, negative translation along the Y-axis, positive translation along the Z-axis, and negative translation along the Z-axis.

[0026] Step 3.3: Based on the point cloud P that the source object point cloud is currently transformed into at step i. i The reward function is constructed based on the Chamfer distance between the target point cloud P' and the target point cloud. Where: η is set to 0.5, S i S i-1 S and S' represent the morphology of the source point cloud at step i and step (i-1), and the state of the target point cloud, respectively. Using the above action space and reward function, a reinforcement learning agent can be trained to generate point cloud registration and fitting trajectories, thereby completing the estimation of the object's 6D pose.

[0027] Step 4: Construct and optimize the kinematic energy function of the joint object, and calculate the pose of the sub-components, specifically including:

[0028] Step 4.1: For the k-th sub-part of the joint object, first initialize the pose of the sub-part with the pose {r0,t0} of the root part.

[0029] Step 4.2: Given the predefined joint information of the sub-component in the point cloud of the source object {u k ,q k}, construct the energy function of the joint based on the kinematic structure of the input joint object: Where: P k For the source point cloud corresponding to this sub-component, P' k For the target point cloud corresponding to this sub-component, N and N' represent P respectively. k and P' k The number of points, T is the transformation function defined based on the kinematic chain of the joint object, T(P) k ,u k ,q k ,θ k ) represents point cloud P k In the joint {u k ,q k The state of} is θ k The point cloud is obtained by time-varying energy functions. Optimizing the energy function above yields the joint {u}. k ,qk} optimal state

[0030] Step 4.3: Determine the status of each joint. By integrating and converting the data into a 6D pose representation, the poses of each sub-component can be obtained. k ,t k}

[0031] Through practical experiments, PointNet++ was used as the encoder to extract features from the input point cloud data, and the final layer dimension was set to 2048. The batch size was set to 32, and the model was trained using the Adam stochastic gradient descent optimizer. State-of-the-art performance was achieved on the PartNet-Mobility dataset, specifically: for the washing machine category, this method achieved a rotation estimation error of 0.9° and 1.1° with only 0.9 seconds of inference time.

[0032]

[0033] In summary, compared with the prior art, the present invention achieves a lower estimation error with an extremely short inference time in the joint object pose estimation task, and can solve the self-occlusion problem in the joint pose estimation task to a certain extent.

[0034] The above-described specific implementations can be partially adjusted by those skilled in the art in different ways without departing from the principles and purpose of the present invention. The scope of protection of the present invention is defined by the claims and is not limited to the above-described specific implementations. All implementation schemes within the scope of the claims are bound by the present invention.

Claims

1. A 6D pose estimation method for jointed objects based on reinforcement learning, characterized in that, The 6D pose estimation problem is transformed into a partial point cloud registration task. The point cloud registration trajectory is used as a sample to train the policy network, which is used to align the source point cloud, i.e., the complete object in the standard space, to the target point cloud, i.e., the partial object observed in the camera space. Then, the trained policy network is used to generate the state for training the agent. The action space and reward function of reinforcement learning are constructed to train the agent. Finally, the kinematic energy function of the joint object is constructed and optimized to calculate the pose of the sub-parts. The policy initialization mentioned above refers to: constructing a PointNet++ structure as a policy network for imitation learning, based on the source object point cloud. and target part point cloud Process and output the serialized pose fitting trajectory ,in: and Representing the first 3D rotation and translation of the steps; The policy network described above uniformly samples the true 6D pose of the object to obtain a real sequential fitting trajectory. Then, the cross-entropy loss function is used to construct the training loss function. To train the policy network, specifically: ,in: Measuring rotational error, Measuring translation error and As a scaling factor, The cross-entropy loss function; The reinforcement learning algorithm described above comprises three main elements: state Action space Reward function , where: state The feature vectors output from the last layer of the PointNet++ architecture represent the state. Design 12 actions to construct the motion space The algorithm consists of six rotation operations: positive X-axis rotation, negative X-axis rotation, positive Y-axis rotation, negative Y-axis rotation, positive Z-axis rotation, and negative Z-axis rotation. The other six operations represent translation operations: positive X-axis translation, negative X-axis translation, positive Y-axis translation, negative Y-axis translation, positive Z-axis translation, and negative Z-axis translation. The reward function is... Based on the The point cloud that the source object point cloud is transformed into at each step. With target point cloud The Chamfer distance between them is designed as follows: ,in: , and These represent the source point cloud at the th Step, First The shape and the state corresponding to the target point cloud.

2. The 6D pose estimation method for jointed objects based on reinforcement learning according to claim 1, characterized in that, The proportionality factor and The values ​​are 0.3 and 0.7 respectively.

3. The 6D pose estimation method for jointed objects based on reinforcement learning according to claim 1, characterized in that, The result of the reward function It is 0.

5.

4. The 6D pose estimation method for jointed objects based on reinforcement learning according to claim 1, characterized in that, The kinematic energy function refers to the function that, when estimating the attitude of a jointed object's sub-components, first uses the attitude of the root component. The poses of each sub-component are initialized. Based on the information of each joint defined in the point cloud of the source object and the kinematic structure of the input joint object, an energy function for the state of each joint is constructed. The energy function is optimized to find the optimal state of the joint. The poses of each sub-component are obtained by integrating the states of each joint.

5. A reinforcement learning-based 6D pose estimation system for jointed objects that implements the method of any one of claims 1-4, characterized in that, include: The system comprises a point cloud feature extraction unit, a reinforcement learning policy initialization unit, a reinforcement agent training unit, and a sub-component pose optimization unit. Specifically: the point cloud feature extraction unit uses the PointNet++ architecture to process point cloud data from each training stage to obtain feature vectors suitable for subsequent processing; the reinforcement learning policy initialization unit utilizes imitation learning to initialize the policy network parameters based on real sequential fitting trajectories; the reinforcement agent training unit uses predefined state, action space, and reward functions to deeply optimize the policy network, training an agent capable of generating point cloud registration trajectories; and the sub-component pose optimization unit optimizes a predefined energy function to calculate the most probable pose of each sub-component of the object.

Images