A training method of a cross-domain segmentation model for visual reinforcement learning generalization

Abstract

The application discloses a training method of a cross-domain segmentation model for visual reinforcement learning generalization, belongs to the technical field of visual reinforcement learning and zero-shot generalization, and aims to alleviate the generalization vulnerability problem of existing visual reinforcement learning algorithms in distribution shift scenarios due to the dependence on non-causal features. First, an RGB channel threshold method is designed to generate a foreground mask from an original image as a mask label for training the cross-domain segmentation model, thereby breaking through the annotation bottleneck. Then, the RGB channel threshold method and the segmentation model are combined to realize self-supervised foreground segmentation, thereby filtering irrelevant interference at the pixel level. Finally, the foreground alignment task and the foreground reconstruction task are coordinated to supplement key background information and strengthen the attention to foreground features, so that a complete and pure state representation is learned. The application can avoid the interference caused by the background from the root cause, and learn representation information conducive to decision-making, thereby improving the zero-shot generalization capability of an intelligent agent.

Description

Technical Field

[0001] This invention belongs to the field of visual reinforcement learning and zero-shot generalization technology, specifically relating to a training method for a cross-domain segmentation model for visual reinforcement learning generalization. Background Technology

[0002] Visual reinforcement learning (RL) directly utilizes image input for policy learning and has achieved breakthroughs in complex tasks such as Atari games and robot manipulation. However, when faced with the inevitable visual distribution shift between training and testing environments, existing methods often experience sharp performance degradation due to over-reliance on non-causal features such as background texture and illumination variations. This generalization fragility mainly stems from the difficulty of visual RL algorithms in effectively decoupling task-related foreground semantics from dynamic background interference during representation learning. Therefore, how to establish cross-domain robust visual representations in complex dynamic scenes remains a core challenge that urgently needs to be addressed.

[0003] To improve the zero-shot generalization ability of visual RL agents, current research mainly focuses on data augmentation techniques and learning robust representations. In data augmentation, mainstream methods expand the distribution boundary of training data by introducing controlled perturbations, thereby enhancing the agent's generalization performance in unknown domains. More refined research attempts to restrict augmentation to specific regions or design task-adaptive augmentation strategies, aiming to improve generalization performance while ensuring that key information is not destroyed. Regarding learning robust representations, research focuses on the representation learning process in visual RL. Some methods learn generalizable representations through pre-trained visual encoders. Other methods design different auxiliary tasks to decouple or highlight key features, thereby obtaining task-relevant and perturbation-invariant representations. However, in reality, due to the pervasive inherent correlation between foreground and background, existing methods often struggle to filter out irrelevant background information at the source.

[0004] However, previous studies have not effectively utilized foreground segmentation, a crucial pixel-level visual prior, despite its significant advantages in object detection and video analysis. This is due to two main reasons: firstly, supervised foreground segmentation models heavily rely on pixel-level annotations (such as the COCO dataset), leading to high annotation costs and pre-training requirements that are difficult to adapt to end-to-end reinforcement learning; secondly, mainstream unsupervised segmentation methods, due to their high computational complexity, cannot meet the real-time requirements of policy updates. It is worth noting that existing work has attempted to extract task-relevant foreground features from pixel inputs using saliency map guidance or image reconstruction methods, but how to construct a cross-domain consistent self-supervised segmentation framework to enhance generalization ability still lacks systematic exploration. Summary of the Invention

[0005] Purpose of the invention: To address the above problems, this invention proposes a training method for a cross-domain segmentation model for generalization in visual reinforcement learning. By using the RGB channel thresholding method to guide the cross-domain segmentation model, it breaks through the annotation bottleneck and avoids the influence of the background from the source. It also coordinates foreground alignment and foreground reconstruction tasks to obtain a complete and clean state representation.

[0006] Technical Solution: To achieve the objectives of this invention, the technical solution adopted is: a training method for a cross-domain segmentation model used for visual reinforcement learning generalization, comprising the following steps:

[0007] Establish a segmentation model network Encoder network Decoder network Policy Network and value network ;

[0008] Establishing a value network Corresponding target value network Target Value Network Structure and Value Network The structure is consistent; the cross-domain segmentation model includes a segmentation model network. and encoder network ;

[0009] The decoder network Policy Network Value Network and target value network Used for training the cross-domain segmentation model.

[0010] Furthermore, training the cross-domain segmentation model includes the following steps:

[0011] Step 1, Update the segmentation model network parameters ;

[0012] Step 2, update encoder network parameters and decoder network parameters ;

[0013] Step 3, update the policy network parameters and value network parameters ;

[0014] Step 4, update the target value network parameters Specifically, it is to Assign values ​​to target value network parameters ;in, Indicates the target value network update rate;

[0015] Repeat steps 1 through 4 until the optimal strategy is obtained.

[0016] Furthermore, the RGB channel thresholding method was used to analyze the original observed image. Generate a foreground mask in the middle;

[0017] The original observation image After the enhancement operation, the enhanced image is obtained. ;

[0018] The enhanced image The foreground mask is used as the mask label in the input image of the segmentation model network. The segmentation model network is trained and its parameters are updated. This enables cross-domain self-supervised segmentation.

[0019] Furthermore, the RGB channel thresholding method was used to analyze the original observed image. Generating a foreground mask involves the following steps: setting minimum and maximum thresholds for the R, G, and B channels; if the original observed image... If the values ​​of all three channels of a pixel are within the corresponding threshold range, then the pixel is determined to be foreground and a foreground mask is generated.

[0020] Furthermore, the minimum and maximum threshold values ​​for the R, G, and B channels are set as follows:

[0021] ;

[0022] ;

[0023] ;

[0024] in This represents the minimum threshold of the R channel. This indicates the maximum threshold of the R channel. This represents the minimum threshold of the G channel. This indicates the maximum threshold value for the G channel. This represents the minimum threshold for channel B. This represents the maximum threshold of channel B.

[0025] Furthermore, the segmentation model network It includes 11 hidden layers, each of which consists of dilated convolutional layers;

[0026] The decoder network It includes 4 hidden layers, each of which consists of 3 convolutional layers and one fully connected layer;

[0027] The encoder network It includes 12 hidden layers, each of which consists of 11 convolutional layers and one fully connected layer;

[0028] Policy Network It includes two hidden layers, and the hidden layers are multilayer perceptron structures with 256 neurons;

[0029] Value Network It includes a multilayer perceptron structure with two hidden layers, each containing 256 neurons.

[0030] Furthermore, the loss function of the segmentation model network is minimized using gradient descent. Update segmentation model network parameters The loss function of the segmentation model network Represented as:

[0031] ;

[0032] ;

[0033]

[0034] in, This represents the loss function of the segmentation model network. The summation symbol is used to represent the summation symbol. This represents the weight of the side output loss of the m-th layer. The weights represent the fusion output loss. Indicates the fusion output loss;

[0035] This represents the side output loss of the m-th layer. These are pixel coordinates, where r represents the x-coordinate and c represents the y-coordinate. Represents the pixel coordinates of the mask label Value, here , Represents the original observation image The c-th channel, This indicates taking the intersection of the three channels. This represents the minimum threshold of the c-th channel. This represents the maximum threshold value for the c-th channel;

[0036] This represents the pixel coordinates of the predicted saliency map generated by the m-th hidden layer of the segmentation model network. value, Let represent the pixel coordinates obtained by overlaying m predicted saliency maps. value;

[0037] Furthermore, the foreground alignment task loss function is minimized using gradient descent. Update encoder network parameters ;

[0038] Minimize the loss function of the foreground reconstruction task using gradient descent. Update decoder network parameters ;

[0039] Foreground alignment task loss function Loss function for foreground reconstruction task They are represented as follows:

[0040] ;

[0041] ;

[0042] in, Expressing expectations, Indicates encoder, Represents the original observed image. Represents the Hadamard product; This indicates the decoder used for foreground reconstruction;

[0043] ;

[0044] B represents a random natural image. This represents the Hadamard product. This represents the probability of allocation.

[0045] Furthermore, utilize the experience pool empirical samples Policy network parameters and value network parameters The update is performed by minimizing the policy network loss function using gradient descent. and the loss function of the value network The loss function is expressed as:

[0046]

[0047]

[0048] in, and Let them represent the state and action at time t, respectively. Represents the adaptive entropy temperature coefficient. express Obedience is based on Action distribution, This represents the reward at time t. This represents the state value function at time t+1.

[0049] Furthermore, segmentation model network Decoder network Encoder network Policy Network and value network The network parameters are initialized randomly.

[0050] Furthermore, target value network parameters The initialization method is as follows: set the corresponding value network parameters Directly assign values ​​to the target value network parameters .

[0051] Beneficial effects: Compared with the prior art, the technical solution of the present invention has the following beneficial technical effects:

[0052] This invention addresses the cross-domain generalization problem in visual reinforcement learning by proposing a training method for a cross-domain segmentation model used for generalization in visual reinforcement learning. (1) The cross-domain segmentation model can stably extract foreground images from complex observation images, thereby fundamentally avoiding interference from irrelevant backgrounds. (2) This method overcomes the labeling bottleneck and meets the high timeliness requirements of reinforcement learning. The foreground mask labels used for training the cross-domain segmentation model are obtained through the RGB channel thresholding method. It avoids the multiple sampling process in previous methods and can directly obtain the optimal action within the action policy support set through guided sampling. (3) This method coordinates foreground pairs with foreground reconstruction tasks to supplement the missing effective background information of pure foregrounds, thereby obtaining a complete and pure state representation. Attached Figure Description

[0053] Figure 1 This is a flowchart of a cross-domain segmentation model used for generalization in visual reinforcement learning. Detailed Implementation

[0054] The technical solution of the present invention will be further described below with reference to the accompanying drawings and embodiments.

[0055] The training method for the cross-domain segmentation model used for visual reinforcement learning generalization described in this invention has the following process: Figure 1 The cross-domain segmentation model training method of this invention includes a cross-domain self-supervised segmentation process and a dual-task representation learning process. First, in the cross-domain self-supervised segmentation process, the original observations are... To construct a complex enhanced image, two enhancement operations, random overlay and random background, are applied. Meanwhile, in the original observation image The method employs RGB channel thresholding to generate a high-confidence foreground mask, which is then used as a mask label for supervision. The mask label Used as a target mask in segmentation model networks for image enhancement. The training proceeds on top of this. Secondly, in the dual-task representation learning process, the dual tasks refer to the foreground alignment task and the foreground reconstruction task, which add effective representation information. The foreground alignment task constrains the original observed image and its corresponding foreground label to pass through the encoder network. The state representations Z and Z obtained respectively * To maintain consistency, the foreground label is composed of a mask label. The image is obtained by multiplying it with the original observed image; the foreground reconstruction task requires extracting the foreground from the enhanced image through an encoder network. and decoder network Constructing the target foreground. During the testing phase, the observed image is processed through the pre-trained cut model network of this invention to obtain a mask. The mask is then multiplied by the observed image to obtain the foreground image, which is then fed into the encoder network. The state representation Z is obtained and used in the subsequent reinforcement learning (RL) process.

[0056] The method includes the following specific steps:

[0057] Step 1: Build the segmentation model network Decoder network Encoder network Policy Network Value Network The network parameters are initialized randomly; these parameters include the segmentation model network parameters. Decoder network parameters Encoder network parameters Policy network parameters and value network parameters ;

[0058] The segmentation model network It includes 11 hidden layers, each of which consists of dilated convolutional layers;

[0059] The decoder network It includes 4 hidden layers, each of which consists of 3 convolutional layers and one fully connected layer;

[0060] The encoder network It includes 12 hidden layers, each of which consists of 11 convolutional layers and one fully connected layer;

[0061] Policy Network It includes two hidden layers, and the hidden layers are multilayer perceptron structures with 256 neurons;

[0062] Value Network It includes a multilayer perceptron structure with two hidden layers, each containing 256 neurons;

[0063] Set RGB thresholds separately ; ; ,in This represents the minimum threshold of the R channel. This indicates the maximum threshold of the R channel. This represents the minimum threshold of the G channel. This indicates the maximum threshold value for the G channel. This represents the minimum threshold for channel B. This represents the maximum threshold of channel B.

[0064] Step 2: Establish the target value network and initialize the target value network parameters.

[0065] Establish the target value network corresponding to the value network ; This represents the target value network parameters corresponding to the target value network; the initialization method for the target value network parameters is: [the method is described in the original value network parameters]. Directly assign values ​​to the target value network parameters .

[0066] Step 3: Update the segmentation model network parameters The specific methods are as follows;

[0067] Using the original observation images in the source domain and corresponding enhanced images For the segmentation model network parameters The update is performed by minimizing the following loss function using gradient descent:

[0068]

[0069]

[0070] in, This represents the side output loss of the m-th layer. These are pixel coordinates, where r represents the x-coordinate and c represents the y-coordinate. Represents the pixel coordinates of the mask label Value, here , Represents the original observation image The c-th channel, This indicates taking the intersection of the three channels. This represents the minimum threshold of the c-th channel. This represents the maximum threshold value for the c-th channel;

[0071] This represents the pixel coordinates of the predicted saliency map generated by the m-th hidden layer of the segmentation model network. Value, final fused graph Obtained by overlaying six intermediate saliency maps;

[0072] This represents the loss function of the segmentation model network. The summation symbol is used to represent the summation symbol. This represents the weight of the output loss at each layer's side. The weights representing the fusion output loss

[0073] Fusion diagram Similarly, the mask label image Calculate binary cross-entropy loss , Indicates the fusion output loss;

[0074] Step 4: Update encoder network parameters and decoder network parameters .

[0075] Using raw observation images in the source domain and corresponding enhanced images Encoder network parameters and decoder network parameters The update is performed by minimizing the following loss function using gradient descent:

[0076]

[0077]

[0078] in, This represents the loss function for the foreground alignment task. Expressing expectations, Indicates encoder, Represents the original observed image. Represents the Hadamard product;

[0079] Loss function for representation and foreground reconstruction tasks This indicates the decoder used for foreground reconstruction. This represents a data-augmented image containing random overlays and random backgrounds. ;

[0080] B represents a random natural image. This represents the Hadamard product. This represents the probability of allocation.

[0081] Step 5: Update the policy network parameters and value network parameters ;

[0082] Utilize the experience pool empirical samples Policy network parameters and value network parameters The update is performed by minimizing the following objective function using gradient descent:

[0083]

[0084]

[0085] in, The loss function of the policy network is represented. and Let them represent the state and action at time t, respectively. Represents the adaptive entropy temperature coefficient. express Obedience is based on Action distribution, Indicating in the value network and target value network Select the minimum value from the output values;

[0086] The loss function represents the value network. This represents the reward at time t. This represents the state value function at time t+1.

[0087] In a robot manipulation task scenario, the state at time t is... This includes the distance between the robotic arm and the target, the robotic arm's speed at time t, and other parameters related to the action. This includes the rotation angle of the robotic arm and the movement speed of the robotic arm.

[0088] Step 6: Update the target value network parameters.

[0089] calculate: ;

[0090] Then, the results of the above calculations are assigned to the target value network parameters. ;in, This represents the target value network update rate.

[0091] Step 7: Repeat steps 3 through 6 until the optimal strategy is obtained.

[0092] Update the parameters of each network and the target value network. The updated cross-domain segmentation model helps the reinforcement learning agent learn a clean and decision-making-friendly state representation.

[0093] This invention relates to a segmentation model for visual reinforcement learning. This segmentation model can avoid interference from irrelevant backgrounds at the source, thereby enhancing the zero-shot generalization ability of reinforcement learning. Specifically, given the dependence of visual reinforcement learning on background texture and illumination, a cross-domain self-supervised segmentation model is introduced. Then, to overcome the annotation bottleneck and meet the high timeliness requirements of reinforcement learning, an RGB channel thresholding method is designed to obtain mask labels for training the segmentation model. Finally, foreground alignment and foreground reconstruction tasks are used in conjunction to supplement the missing effective background information of the pure foreground, thereby obtaining a complete and pure state representation.

[0094] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the principle of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. A method for training a cross-domain segmentation model, the method comprising: The method comprises the following steps: establishing a segmentation model network , an encoder network , a decoder network , a policy network , and a value network ; Establishing a value network The corresponding target value network The structure of the target value network is consistent with the structure of the value network ​ The cross-domain segmentation model comprises a segmentation model network and an encoder network ; The decoder network , a policy network , a value network , and a target value network For training the cross-domain segmentation model, specifically comprising the following steps: Step 1, updating the segmentation model network parameters ; The foreground mask is generated from the original observation image using an RGB channel thresholding method. The foreground mask is generated from the original observation image using an RGB channel thresholding method. performing an enhancement operation on the original observation image performing an enhancement operation on the original observation image ; the enhanced image as an input image to the segmentation model network the foreground mask as a mask label the segmentation model network is trained by minimizing a segmentation model network loss function updates the segmentation model network parameters the segmentation model network loss function is represented as: ； ； wherein, denotes a segmentation model network loss function, denotes a summation symbol, denotes a weight of the m-th layer side edge output loss, denotes a weight of the fusion output loss, denotes a fusion output loss; represents the m-th layer side edge output loss, is the pixel coordinate, r represents the horizontal coordinate, and c represents the vertical coordinate, represents the pixel coordinate of the mask label value, here , represents the original observation image of the c-th channel, represents the intersection of the three channels, represents the minimum threshold value of the c-th channel, represents the maximum threshold value of the c-th channel; pixel coordinates of the prediction saliency map generated by the mth hidden layer of the segmentation model network value, pixel coordinates obtained by superimposing the m prediction saliency maps value; Step 2, minimizing the foreground alignment task loss function by gradient descent updating the encoder network parameters ; minimizing the foreground reconstruction task loss function by gradient descent updating the decoder network parameters ; foreground alignment task loss function and foreground reconstruction task loss function are represented as follows, respectively: ； ； wherein, denotes expectation, denotes encoder, denotes original observed image, denotes Hadamard product; denotes decoder for foreground reconstruction; Step 3, updating the policy network parameters and the value network parameters using the experience samples in the experience pool ​​​​​​ wherein, and s(t) and a(t) represent the state and action at time t, respectively, represents an adaptive entropy temperature coefficient, represents obeys an action distribution based on represents represents the reward at time t, represents the state value function at time t+1; Step 4, updating the target value network parameters , specifically assigning to the target value network parameters ; wherein represents the target value network update rate; Repeat steps 1 to 4 until the optimal strategy is obtained. 2.The method of claim 1, wherein, The original observation image is obtained After the enhancement operation, an enhanced image is obtained is expressed as: B represents a random natural image, denotes the Hadamard product, denotes the assignment probability. 3.The method of claim 1, wherein, Using RGB channel thresholding method to extract from the original observation image Generating a foreground mask involves the following steps: Set minimum threshold and maximum threshold of R, G and B channels; If the values of the three channels of the pixel point in the original observation image are all within the corresponding threshold range, the pixel point is determined as foreground, and a foreground mask is generated.

4. The method of claim 3, wherein the training of the cross-domain segmentation model is performed by, The minimum threshold and the maximum threshold of the R, G and B channels are respectively: ； ； ； wherein represents the R channel minimum threshold, represents the R channel maximum threshold, represents the G channel minimum threshold, represents the G channel maximum threshold, represents the B channel minimum threshold, represents the B channel maximum threshold.

5. The method of claim 1, wherein, The segmentation model network includes 11 hidden layers, each of which is composed of an expanded convolutional layer; The decoder network includes 4 hidden layers, each hidden layer including 3 convolutional layers and one fully connected layer; The encoder network includes 12 hidden layers, each hidden layer including 11 convolutional layers and one fully connected layer; Policy network includes two hidden layers, and the hidden layers are multilayer perceptron structures with 256 neurons Value network includes two hidden layers, and the hidden layers are multi-layer perceptron structures with 256 neurons.

6. The method of Claim 1, wherein, Split model network Decoder network Encoder network Policy network Value network The initialization of the network parameters of the split model network, the decoder network, the encoder network, the policy network, and the value network is random initialization. target value network parameter The initialization method is to directly assign the corresponding value network parameter to the target value network parameter .

Images