Strategy network training method and system based on future latent variable alignment and flow matching
Patent Information
- Application Number
- CN202611108243.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2026-07-24
- Publication Date
- 2026-10-30
AI Technical Summary
[0006]为了克服现有机器人策略学习方法要么因缺乏对未来状态的推理能力而难以应对长时域任务，要么因依赖像素级未来图像生成而导致计算开销过大且优化目标与机器人控制需求错位的问题，本发明提供了一种基于未来潜变量对齐与流匹配的策略网络训练方法及系统
[0037]1、本发明通过在流匹配策略网络中增设可学习的未来Token，并在网络中间层将其输出特征与真实未来观测嵌入进行对齐，使策略网络能够将未来状态信息编码于未来Token之中，在动作生成时隐式地感知当前动作将导致的未来状态变化，从而在长时域任务中能够基于对未来状态的预判进行前瞻性决策，克服了传统反应式策略因缺乏规划能力而导致的累积误差问题；
Smart Images
Figure CN123000784A_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of robot control technology, and more specifically, to a method and system for training policy networks based on future latent variable alignment and flow matching. Background Technology
[0002] Learning-based robot operation strategies are one of the core technologies in the field of embodied intelligence. Currently, existing methods for learning robot operation strategies mainly fall into two categories.
[0003] The first category is policy network methods based on imitation learning. Behavioral cloning achieves robot control by establishing a direct mapping from current observations to actions, but this method is prone to getting trapped in local optima due to the multimodal distribution of training data. In recent years, diffusion strategies and flow matching strategies have emerged, using generative models to fit the data distribution, which has improved performance on multimodal tasks to some extent. However, these methods are inherently reactive, relying solely on the current observation input to generate actions, lacking explicit reasoning ability regarding future changes in the environmental state. When facing long-term tasks requiring multi-step decision-making, these methods suffer from low success rates due to a lack of forward-looking planning.
[0004] The second category is control methods based on world models. These methods attempt to combine video generation models with policy networks, generating future image sequences to "imagine" the consequences of actions, thereby aiding decision-making. However, these methods require high-fidelity reconstruction at the pixel level, demanding enormous computational power to generate high-resolution future image sequences. Inference times typically reach hundreds of milliseconds or even seconds, making it difficult to meet the demands of high-frequency, real-time robot control. Furthermore, the optimization goal of visual generation models is to improve image perception quality, while the core focus of robot control is on object position and geometric state; their optimization objectives are fundamentally misaligned.
[0005] Therefore, there is an urgent need to provide a solution to the above problems. Summary of the Invention
[0006] To overcome the problems of existing robot policy learning methods, which either lack the ability to reason about future states and are therefore unable to handle long-term tasks, or rely on pixel-level future image generation, resulting in excessive computational overhead and a misalignment between optimization objectives and robot control requirements, this invention provides a policy network training method and system based on future latent variable alignment and flow matching.
[0007] The technical solution of this invention is as follows:
[0008] A method for training a policy network based on future latent variable alignment and flow matching includes the following steps:
[0009] Training data containing the current observation image, the future observation image, and the real action sequence is acquired. Based on the pre-trained action-aware embedding model, the current observation image and the future observation image are mapped to obtain the current observation embedding and the target future feature. The real action sequence is denoised to generate a noisy action sequence, and the noisy action sequence is embedded to obtain the noisy action embedding.
[0010] A flow matching policy network is constructed based on the Transformer architecture, and a learnable future token is added to the policy network. The current observation embedding, the noisy action embedding, and the future token are concatenated into an input sequence, which is then input into the policy network for forward propagation to obtain a predicted action sequence. The output features corresponding to the future token are extracted in the middle layer of the policy network to obtain the predicted future features.
[0011] Calculate the similarity loss between the predicted future feature and the target future feature; calculate the action flow matching loss based on the error between the real action sequence and the predicted action sequence; construct a hybrid loss function based on the similarity loss and the action flow matching loss; and update the parameters of the policy network based on the hybrid loss function.
[0012] Preferably, the pre-training process of the action-aware embedding model includes:
[0013] The sample image is input into a visual encoder for encoding to obtain visual features; the visual features are then input into a feature compression module for compression to obtain a compact embedding vector.
[0014] The compact embedded vector input stream matching decoder is used for action prediction to obtain the predicted action sequence;
[0015] The pre-training loss is calculated based on the error between the predicted action sequence and the sample action sequence, and the parameters of the visual encoder and the feature compression module are updated based on the pre-training loss.
[0016] Preferably, the feature compression module includes M learnable Query Tokens. The process of inputting the visual features into the feature compression module for compression processing to obtain a compact embedding vector includes:
[0017] The visual features and the M Query Tokens are input into the feature compression module. Based on the cross-attention calculation between the Query Tokens and the visual features, the visual features are compressed into M embedding vectors to obtain the compact embedding vector.
[0018] Preferably, the process of generating a noisy action sequence by adding noise to the real action sequence includes:
[0019] Obtain a preset number of diffusion steps and a noise scheduling function, and determine the noise intensity of the current step based on the noise scheduling function; add Gaussian noise to the real action sequence based on the noise intensity to generate a noisy action sequence.
[0020] Preferably, the process of extracting the output features corresponding to the future token in the intermediate layer of the policy network to obtain the predicted future features includes:
[0021] Obtain the total number of layers L of the policy network. During the forward propagation of the policy network, extract the output features corresponding to the future token at the L / 2 layer to obtain the predicted future features.
[0022] Preferably, the process of calculating the similarity loss between the predicted future feature and the target future feature includes:
[0023] Calculate the cosine similarity between the predicted future feature and the target future feature to obtain a cosine similarity value; calculate the similarity loss based on the cosine similarity value, and the similarity loss is negatively correlated with the cosine similarity value.
[0024] Preferably, the expression for the hybrid loss function is:
[0025] ;
[0026] in, For mixed loss, The loss is used to match the action flow. For the aforementioned similarity loss, The weighting coefficients are preset and are used to adjust the contribution ratio of the similarity loss in the hybrid loss function.
[0027] A second aspect of the present invention also provides a policy network training system based on future latent variable alignment and flow matching, comprising a data preprocessing module, a policy network forward propagation module, and a policy network parameter update module, wherein:
[0028] The data preprocessing module is used to acquire training data containing the current observation image, the future observation image, and the real action sequence. Based on the pre-trained action-aware embedding model, the current observation image and the future observation image are mapped to obtain the current observation embedding and the target future feature. The real action sequence is denoised to generate a noisy action sequence, and the noisy action sequence is embedded to obtain the noisy action embedding.
[0029] The policy network forward propagation module is used to construct a flow matching policy network based on the Transformer architecture, and add a learnable future token to the policy network. The current observation embedding, the noisy action embedding and the future token are concatenated into an input sequence, which is then input into the policy network for forward propagation to obtain a predicted action sequence. The output features corresponding to the future token are extracted in the middle layer of the policy network to obtain the predicted future features.
[0030] The policy network parameter update module is used to calculate the similarity loss between the predicted future feature and the target future feature, calculate the action flow matching loss based on the error between the real action sequence and the predicted action sequence, construct a hybrid loss function based on the similarity loss and the action flow matching loss, and update the parameters of the policy network based on the hybrid loss function.
[0031] A third aspect of the present invention also provides an electronic device, the electronic device comprising:
[0032] At least one processor; and,
[0033] A memory communicatively connected to the at least one processor; wherein,
[0034] The memory stores a computer program that can be executed by the at least one processor, which enables the at least one processor to perform a policy network training method based on future latent variable alignment and flow matching as described above.
[0035] A fourth aspect of the present invention also provides a computer-readable storage medium, including a storage data area and a storage program area, wherein the storage data area stores created data and the storage program area stores a computer program; wherein, when the computer program is executed by a processor, it implements a policy network training method based on future latent variable alignment and flow matching as described above.
[0036] The beneficial effects of the policy network training method and system based on future latent variable alignment and flow matching according to the above scheme are as follows:
[0037] 1. This invention adds a learnable future token to the flow matching policy network and aligns its output features with the embedding of real future observations in the intermediate layer of the network. This enables the policy network to encode future state information in the future token and implicitly perceive the future state changes that the current action will cause when the action is generated. Thus, in long-term tasks, it can make forward-looking decisions based on the prediction of future states, overcoming the problem of accumulated error caused by the lack of planning ability in traditional reactive policies.
[0038] 2. The strategy network of the present invention directly outputs the predicted future features in the form of embedding and performs similarity calculation with the real future features in the form of embedding. It does not require the reconstruction of future observation images at the pixel level, thus avoiding the huge computing power consumption and inference delay caused by high-resolution image generation, and can meet the needs of high-frequency real-time control of robots.
[0039] 3. Existing video generation models optimize image perception quality, which is misaligned with the core focus of robot control on object position and geometric state. This invention uses the accuracy of future state prediction in feature space as the alignment target, so that the representation learned by the model directly serves the action generation task, avoiding the decrease in learning efficiency and performance loss caused by the misalignment of optimization targets. Attached Figure Description
[0040] To more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.
[0041] Figure 1 This is a schematic diagram of the strategy network training method provided by the present invention;
[0042] Figure 2 This is a schematic diagram of the structure of an electronic device provided in an embodiment of the present invention. Detailed Implementation
[0043] To make the technical problems to be solved, the technical solutions, and the beneficial effects of the present invention clearer, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.
[0044] It should be noted that when a component is referred to as "fixed," "set," or "connected" to another component, it may be located directly or indirectly on that other component. The terms "upper," "lower," "left," "right," "front," "rear," "vertical," "horizontal," "top," "bottom," "inner," and "outer," etc., indicate the orientation or position based on the accompanying drawings, and are for ease of description only, and should not be construed as limiting the technical solution. The terms "first," "second," etc., are used for ease of description only and should not be construed as indicating or implying relative importance or implicitly specifying the number of technical features. "Many" means two or more, unless otherwise explicitly specified. "Several" means one or more, unless otherwise explicitly specified.
[0045] This invention provides a method for training policy networks based on future latent variable alignment and flow matching. See [link to relevant documentation]. Figure 1 This includes the following steps:
[0046] S1. Obtain training data containing the current observation image, the future observation image, and the real action sequence. Based on the pre-trained action-aware embedding model, perform mapping processing on the current observation image and the future observation image respectively to obtain the current observation embedding and the target future feature. Add noise to the real action sequence to generate a noisy action sequence, and perform embedding processing on the noisy action sequence to obtain the noisy action embedding.
[0047] S2. Construct a stream matching policy network based on the Transformer architecture, and add a learnable future token to the policy network. Concatenate the current observation embedding, the noisy action embedding, and the future token into an input sequence, input it into the policy network for forward propagation, obtain the predicted action sequence, and extract the output features corresponding to the future token in the middle layer of the policy network to obtain the predicted future features.
[0048] S3. Calculate the similarity loss between the predicted future feature and the target future feature. Based on the error between the real action sequence and the predicted action sequence, calculate the action flow matching loss. Construct a hybrid loss function based on the similarity loss and the action flow matching loss. Update the parameters of the policy network based on the hybrid loss function.
[0049] In some embodiments, before performing step S1, a pre-trained action-aware embedding model is further included, comprising: acquiring sample images and corresponding sample action sequences, inputting the sample images into a visual encoder for encoding processing to obtain visual features, and then inputting the visual features into a feature compression module for compression processing to obtain a compact embedding vector.
[0050] Specifically, the feature compression module includes M learnable Query Tokens. During compression, the visual features and the M Query Tokens are input into the feature compression module. Through cross-attention calculation between the Query Tokens and the visual features, the visual features are compressed into M embedding vectors, resulting in compact embedding vectors. Preferably, M is 32, and the visual encoder adopts the SigLIP-2 architecture.
[0051] Furthermore, after obtaining the compact embedding vector, the compact embedding vector is input into the stream matching decoder for action prediction to obtain the predicted action sequence. The pre-training loss is calculated based on the error between the predicted action sequence and the sample action sequence, and the parameters of the visual encoder and feature compression module are updated based on this pre-training loss. After pre-training, the stream matching decoder is discarded, and only the visual encoder and feature compression module are retained to form the action-aware embedding model, which is used for feature extraction in the subsequent policy network training stage.
[0052] In some embodiments, the process of performing step S1 includes:
[0053] S1.1 Obtain training data containing the current observation image, the future observation image, and the real action sequence;
[0054] S1.2. Based on the pre-trained action-aware embedding model, the current observation image and the future observation image are mapped to obtain the current observation embedding and the target future features.
[0055] S1.3. Add noise to the real action sequence to generate a noisy action sequence;
[0056] S1.4. Embed the noisy action sequence to obtain the noisy action embedding.
[0057] Specifically, during the execution of step S1.3, a preset diffusion step number and noise scheduling function are obtained. Based on the noise scheduling function, the noise intensity of the current step is determined. Gaussian noise is added to the real action sequence according to the noise intensity to generate a noisy action sequence.
[0058] Furthermore, during the execution of step S1.4, the embedding processing of the noisy action sequence is completed by the input embedding layer of the policy network, which maps the noisy action sequence to a unified vector dimension.
[0059] In some embodiments, the training data also includes video data without action labels, which contains the current observation image and the future observation image. For the video data without action labels, step S1.2 is also performed to perform mapping processing to obtain the current observation embedding and the target future features, which are used only to calculate the similarity loss in subsequent steps and do not participate in the calculation of the action flow matching loss.
[0060] In some embodiments, the process of performing step S2 includes:
[0061] S2.1. Construct a stream matching policy network based on the Transformer architecture, and add a learnable future token to the policy network;
[0062] S2.2, Concatenate the current observation embedding, the noisy action embedding, and the future token into an input sequence;
[0063] S2.3. Input the input sequence into the policy network for forward propagation to obtain the predicted action sequence;
[0064] S2.4 Extract the output features corresponding to future tokens in the intermediate layer of the policy network to obtain the predicted future features.
[0065] Specifically, during step S2.1, the policy network is constructed using a Diffused Transformer (DiT) architecture. The policy network comprises an input embedding layer, a feature processing backbone consisting of stacked L-layer Transformer decoders, and an output projection layer. The input embedding layer maps the current observation embedding, the noisy action embedding, and the future token to a unified vector dimension, and then concatenates them along the sequence dimension. Each layer of the L-layer Transformer decoder contains a multi-head self-attention sublayer and a position-wise feedforward network sublayer.
[0066] Furthermore, the learnable future token is a set of trainable parameter vectors that participate in attention calculation together with other tokens in the input sequence. During the execution of steps S2.3 and S2.4, after the input sequence is processed by L layers of Transformer decoders stacked sequentially, the future token is processed by the first L / 2 layers of Transformer decoders. The output state at its corresponding position fuses the current observation information and action information. The output features of the future token are extracted from this layer to obtain the predicted future features. The hidden state output by the Lth layer is mapped by the output projection layer to obtain the predicted action sequence.
[0067] In some embodiments, the process of performing step S3 includes:
[0068] S3.1 Calculate the similarity loss between the predicted future features and the target future features;
[0069] S3.2 Calculate the action flow matching loss based on the error between the real action sequence and the predicted action sequence;
[0070] S3.3 Construct a hybrid loss function based on similarity loss and action flow matching loss;
[0071] S3.4 Update the parameters of the policy network based on the hybrid loss function.
[0072] Specifically, during step S3.1, the cosine similarity between the predicted future feature and the target future feature is calculated to obtain a cosine similarity value. Based on the cosine similarity value, a similarity loss is calculated, and this similarity loss is negatively correlated with the cosine similarity value. That is, the larger the cosine similarity value (the closer the predicted future feature is to the target future feature), the smaller the similarity loss; conversely, the smaller the cosine similarity value, the larger the similarity loss.
[0073] Specifically, during the execution of step S3.3, the expression for the hybrid loss function is:
[0074] ;
[0075] in, For mixed loss, The loss is used to match the action flow. For the aforementioned similarity loss, The weighting coefficients are preset and are used to adjust the contribution ratio of the similarity loss in the hybrid loss function.
[0076] Furthermore, during step S3.4, the total loss value is calculated based on the hybrid loss function, and the parameters of the policy network are updated through backpropagation.
[0077] In a second aspect of the invention, a policy network training system based on future latent variable alignment and flow matching is also provided, characterized in that it includes a data preprocessing module, a policy network forward propagation module, and a policy network parameter update module, wherein:
[0078] The data preprocessing module is used to acquire training data containing the current observation image, the future observation image, and the real action sequence. Based on the pre-trained action-aware embedding model, the current observation image and the future observation image are mapped to obtain the current observation embedding and the target future feature. The real action sequence is denoised to generate a noisy action sequence, and the noisy action sequence is embedded to obtain the noisy action embedding.
[0079] The policy network forward propagation module is used to construct a flow matching policy network based on the Transformer architecture, and add a learnable future token to the policy network. The current observation embedding, the noisy action embedding and the future token are concatenated into an input sequence, which is then input into the policy network for forward propagation to obtain a predicted action sequence. The output features corresponding to the future token are extracted in the middle layer of the policy network to obtain the predicted future features.
[0080] The policy network parameter update module is used to calculate the similarity loss between the predicted future feature and the target future feature, calculate the action flow matching loss based on the error between the real action sequence and the predicted action sequence, construct a hybrid loss function based on the similarity loss and the action flow matching loss, and update the parameters of the policy network based on the hybrid loss function.
[0081] In a third aspect of the invention, an electronic device is also provided, see [link to relevant documentation]. Figure 2 The electronic device may include a processor 10, a memory 11, a communication bus 12, and a communication interface 13, and may also include a computer program stored in the memory 11 and executable on the processor 10; wherein the computer program is executed by the at least one processor 10 to enable the at least one processor 10 to execute a policy network training method based on future latent variable alignment and flow matching as described above.
[0082] In a fourth aspect of the invention, a computer-readable storage medium is also provided, including a storage data area and a storage program area, wherein the storage data area stores created data and the storage program area stores a computer program; wherein, when the computer program is executed by a processor, it implements a policy network training method based on future latent variable alignment and flow matching as described above.
[0083] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the protection scope of the present invention.
Claims
1. A policy network training method based on future latent variable alignment and flow matching, characterized in that, Includes the following steps: Training data containing the current observation image, the future observation image, and the real action sequence is acquired. Based on the pre-trained action-aware embedding model, the current observation image and the future observation image are mapped to obtain the current observation embedding and the target future feature. The real action sequence is denoised to generate a noisy action sequence, and the noisy action sequence is embedded to obtain the noisy action embedding. A flow matching policy network is constructed based on the Transformer architecture, and a learnable future token is added to the policy network. The current observation embedding, the noisy action embedding, and the future token are concatenated into an input sequence, which is then input into the policy network for forward propagation to obtain a predicted action sequence. The output features corresponding to the future token are extracted in the middle layer of the policy network to obtain the predicted future features. Calculate the similarity loss between the predicted future feature and the target future feature; calculate the action flow matching loss based on the error between the real action sequence and the predicted action sequence; construct a hybrid loss function based on the similarity loss and the action flow matching loss; and update the parameters of the policy network based on the hybrid loss function.
2. The policy network training method based on future latent variable alignment and flow matching as described in claim 1, characterized in that, The pre-training process of the action-aware embedding model includes: The sample image is input into a visual encoder for encoding to obtain visual features; the visual features are then input into a feature compression module for compression to obtain a compact embedding vector. The compact embedded vector input stream matching decoder is used for action prediction to obtain the predicted action sequence; The pre-training loss is calculated based on the error between the predicted action sequence and the sample action sequence, and the parameters of the visual encoder and the feature compression module are updated based on the pre-training loss.
3. The policy network training method based on future latent variable alignment and flow matching as described in claim 2, characterized in that, The feature compression module contains M learnable Query Tokens. The process of inputting the visual features into the feature compression module for compression to obtain a compact embedding vector includes: The visual features and the M Query Tokens are input into the feature compression module. Based on the cross-attention calculation between the Query Tokens and the visual features, the visual features are compressed into M embedding vectors to obtain the compact embedding vector.
4. The policy network training method based on future latent variable alignment and flow matching as described in claim 1, characterized in that, The process of generating a noisy action sequence by adding noise to the real action sequence includes: Obtain a preset number of diffusion steps and a noise scheduling function, and determine the noise intensity of the current step based on the noise scheduling function; add Gaussian noise to the real action sequence based on the noise intensity to generate a noisy action sequence.
5. The policy network training method based on future latent variable alignment and flow matching as described in claim 1, characterized in that, The process of extracting the output features corresponding to the future token in the intermediate layer of the policy network to obtain the predicted future features includes: Obtain the total number of layers L of the policy network. During the forward propagation of the policy network, extract the output features corresponding to the future token at the L / 2 layer to obtain the predicted future features.
6. The policy network training method based on future latent variable alignment and flow matching as described in claim 1, characterized in that, The process of calculating the similarity loss between the predicted future feature and the target future feature includes: Calculate the cosine similarity between the predicted future feature and the target future feature to obtain a cosine similarity value; calculate the similarity loss based on the cosine similarity value, and the similarity loss is negatively correlated with the cosine similarity value.
7. The policy network training method based on future latent variable alignment and flow matching as described in claim 1, characterized in that, The expression for the hybrid loss function is: ； in, For mixed loss, The loss is used to match the action flow. For the aforementioned similarity loss, The weighting coefficients are preset and are used to adjust the contribution ratio of the similarity loss in the hybrid loss function.
8. A policy network training system based on future latent variable alignment and flow matching, characterized in that, It includes a data preprocessing module, a policy network forward propagation module, and a policy network parameter update module, wherein: The data preprocessing module is used to acquire training data containing the current observation image, the future observation image, and the real action sequence. Based on the pre-trained action-aware embedding model, the current observation image and the future observation image are mapped to obtain the current observation embedding and the target future feature. The real action sequence is denoised to generate a noisy action sequence, and the noisy action sequence is embedded to obtain the noisy action embedding. The policy network forward propagation module is used to construct a flow matching policy network based on the Transformer architecture, and add a learnable future token to the policy network. The current observation embedding, the noisy action embedding and the future token are concatenated into an input sequence, which is then input into the policy network for forward propagation to obtain a predicted action sequence. The output features corresponding to the future token are extracted in the middle layer of the policy network to obtain the predicted future features. The policy network parameter update module is used to calculate the similarity loss between the predicted future feature and the target future feature, calculate the action flow matching loss based on the error between the real action sequence and the predicted action sequence, construct a hybrid loss function based on the similarity loss and the action flow matching loss, and update the parameters of the policy network based on the hybrid loss function.
9. An electronic device, characterized in that, The electronic device includes: At least one processor; and, A memory communicatively connected to the at least one processor; wherein, The memory stores a computer program that can be executed by the at least one processor, the computer program being executed by the at least one processor to enable the at least one processor to perform a policy network training method based on future latent variable alignment and flow matching as described in any one of claims 1 to 7.
10. A computer-readable storage medium, comprising a data storage area and a program storage area, wherein the data storage area stores created data and the program storage area stores a computer program; wherein, When the computer program is executed by the processor, it implements a policy network training method based on future latent variable alignment and flow matching as described in any one of claims 1 to 7.