World model-based closed-loop policy pre-exploration action prediction method and device

Abstract

This invention discloses a closed-loop policy prediction method and apparatus based on a world model, relating to the field of action prediction. The method includes: during iterative prediction at the current time step, using a pre-trained world model, the core hidden state of the current time step is obtained based on the visual observation frame and action command of the current time step, as well as the core hidden state of the previous time step. The core hidden state of the current time step is then input into a trained policy model to obtain the action command for the next time step. The pre-trained world model then generates the visual observation frame for the next time step based on the action command for the next time step and the core hidden state of the current time step. This iterative prediction process is repeated to obtain action commands for different time steps, which are then sent to the machine to execute the corresponding actions. This invention addresses the problem of existing world models causing inaccurate action predictions due to the predicted state deviating from the actual physical state.

Description

Technical Field

[0001] This invention relates to the field of action prediction, and more specifically to a closed-loop strategy-based action prediction method and apparatus based on a world model. Background Technology

[0002] In recent years, with the rapid development of intelligent manufacturing, warehousing and logistics, and home service robots, autonomous operation of robotic arms has become a key to industrial upgrading. However, large-scale real-machine interaction training methods are unsustainable. Closed-loop policy propagation prediction based on world models enables machines to complete training in virtual interactive environments, greatly reducing costs and representing an important path to large-scale deployment. However, existing world-model-based methods are prone to prediction illusions during closed-loop policy propagation, making it difficult to accurately depict the state evolution patterns in real-world environments. Specifically, on the one hand, during single-step prediction, coordination misalignment can easily occur between visual representation and action response, causing the current predicted state to deviate from the actual physical state; on the other hand, during long-term time-series extrapolation, local prediction errors gradually accumulate and amplify, further inducing false success signals, interfering with gradient updates during policy prediction, and thus affecting the prediction results of closed-loop policy propagation. Summary of the Invention

[0003] The purpose of this application is to address the technical problem that existing world models cause inaccurate predictions of policy model actions due to deviations between predicted states and actual physical states. This application proposes a closed-loop policy prediction method and apparatus based on a world model.

[0004] In a first aspect, the present invention provides a method for predicting forward actions of a closed-loop policy based on a world model, comprising the following steps:

[0005] Construct a policy model and a pre-trained world model. Based on the pre-trained world model, use a closed-loop prediction process to train the policy model to obtain the trained policy model.

[0006] The system acquires the visual observation frame at the first time step and performs an iterative prediction process for the first time step by combining the action command at the first time step with a value of zero and the core hidden state at the zeroth time step. During the iterative prediction process at the current time step, the system uses a pre-trained world model to obtain the core hidden state at the current time step based on the visual observation frame, action command, and core hidden state at the previous time step. The system inputs the core hidden state at the current time step into a trained policy model to obtain the action command for the next time step. The system then uses the pre-trained world model to generate the visual observation frame for the next time step based on the action command for the next time step and the core hidden state at the current time step. The system repeats the above iterative prediction process to obtain action commands for different time steps and sends them to the machine to execute the corresponding actions.

[0007] As a preferred embodiment, the pre-trained world model includes a prediction branch, a training branch, and a reconstruction branch. The prediction branch includes a visual encoder, a first action encoder, a first adaptive normalization layer, and a first cross-attention module.

[0008] The visual observation frame and action command of the current time step, along with the core hidden state of the previous time step, are input into the prediction branch to obtain the core hidden state of the current time step. The specific process is as follows:

[0009] The visual observation frame and motion command of the current time step are input into the visual encoder and the first motion encoder, respectively, to obtain the visual embedding and motion embedding of the current time step.

[0010] The visual embedding of the current time step and the core hidden state of the previous time step are input into the first adaptive normalization layer. The visual embedding of the current time step is used as the input feature, and the core hidden state of the previous time step is used as the conditional information. The first intermediate feature of the current time step is modulated as shown in the following equation:

[0011] ;

[0012] Where t represents the current time step, and t-1 represents the previous time step. The visual embedding representing the current time step. This indicates the core hidden state of the previous time step. and represents the mean and standard deviation of the core hidden state at the previous time step, respectively; and These represent the scaling mapping function and the offset mapping function, respectively, both of which can be implemented by independent fully connected layers; and These represent the scaling vector and the offset vector, respectively. For element-wise multiplication operator, This represents the first intermediate feature at the current time step;

[0013] The action embedding and the first intermediate feature of the current time step are input into the first cross-attention module. The action embedding of the current time step is used to obtain the key matrix and value matrix through linear projection. The first intermediate feature of the current time step is used to obtain the query matrix through linear projection. The core hidden state of the current time step is calculated through the cross-attention mechanism.

[0014] Preferably, the training branch includes a state deviation measurement module, a keyframe retrieval module, and a state reset module; the state deviation measurement module includes a second adaptive normalization layer, a multilayer perceptron, a first decoder, and a linear projection layer.

[0015] The training process of the policy model includes:

[0016] First, execute the closed-loop prediction process. The steps of the closed-loop prediction process are as follows:

[0017] The visual observation frame and action command of the current time step in the training sample, as well as the core hidden state of the previous time step, are passed through the prediction branch to obtain the core hidden state of the current time step.

[0018] The core hidden state of the previous time step, the core hidden state of the current time step, and the action command from the training samples are input into the training branch. First, they pass through the state bias measurement module. Then, the core hidden state of the previous time step and the action command from the current time step are input into the second adaptive normalization layer. The core hidden state of the previous time step is used as the input feature, and the action command from the current time step is used as the conditional information to modulate and obtain the second intermediate feature of the current time step. The second intermediate feature of the current time step is then input into the multilayer perceptron to map and obtain the reference state of the current time step.

[0019] The core hidden state at the current time step is passed sequentially through the first decoder and the linear projection layer to map the predicted state at the current time step.

[0020] Calculate the deviation between the predicted state and the reference state at the current time step; in response to the determination that the deviation values ​​of N consecutive time steps t-N+1, t-N+2, ..., t exceed the deviation threshold, trigger the state adjustment process. The current time step is used as the time step in which the strategy's forward prediction fails, and the reference state of the current time step is used as the reference state of the time step in which the strategy's forward prediction fails.

[0021] During the state adjustment process The time step where the strategy-predicted failure occurred and its reference state are input into the keyframe retrieval module. The prediction states of the k time steps preceding the time step where the strategy-predicted failure occurred are used to form a neighboring candidate window set for the time step where the strategy-predicted failure occurred. The prediction state of the optimal keyframe is retrieved from this neighboring candidate window set, as shown in the following formula:

[0022] ;

[0023] in, This represents the set of neighboring candidate windows at the time step where the strategy's prediction failed. The index i represents the prediction state in the set of neighboring candidate windows at the time step where the strategy prediction failed, where i = tk, t-k+1, ..., t-1; This represents the prediction state of the optimal keyframe. The independent variable represents the value at which the function reaches its minimum. This indicates the reference state at the time step in which the strategy's prediction failed. This is the timing penalty coefficient. For window length, Represents the L2 norm;

[0024] The prediction state of the optimal keyframe is input into the state reset module, the core hidden state of the time step corresponding to the prediction state of the optimal keyframe is extracted, and it is forcibly overwritten on the core hidden state of the current time step of the time step that failed as the strategy forward prediction.

[0025] The closed-loop prediction process is repeated, and the core hidden state of the current time step, which is the time step where the strategy foreshadowing prediction fails, is adjusted to obtain the core hidden state of each time step.

[0026] The success probability of the task at each time step is obtained by mapping the core hidden state at each time step, as shown in the following formula:

[0027] ;

[0028] in, and These represent the weight matrix and the bias term, respectively. This indicates the probability of task success at the current time step;

[0029] The time step with a success probability greater than the success determination threshold is set as the success time step. The visual observation frames of the success time step and the time steps before it are taken as valid frames and the corresponding mask is set to 1. The visual observation frames of the time steps after the success time step are taken as invalid frames and the corresponding mask is set to 0.

[0030] The total loss function and its gradient used during the training of the strategy model are calculated using the following formula:

[0031] ;

[0032] ;

[0033] ;

[0034] in, Represents the total loss function. This represents the cumulative discounted reward from the successful time step (l) to the successful time step T, where j represents the time step variable, j = l, l+1, ..., T. j-l Represents time step variables j With one of the time steps l difference, Indicates the discount factor. A mask representing the successful time step and one of the time steps l preceding it. This represents the core hidden state of time step l, which is one of the time steps preceding the successful time step. After being input into the policy model, the action command for the next time step is output. The conditional probability, Indicates the parameter to be optimized is The strategy model; This represents the probability of task success for time step variable j; The parameters to be optimized in the strategy model Gradient operator for variables;

[0035] The parameters to be optimized in the policy model are updated using the gradient of the total loss function, resulting in a trained policy model.

[0036] Preferably, both the first and second adaptive normalization layers are adaptive normalization layers; the visual encoder adopts the Vision Transformer structure, and the deviation value is the square of the Euclidean distance.

[0037] As a preferred embodiment, the reconstruction branch includes a second action encoder, a second cross-attention module, a conditional diffusion model, and a second decoder.

[0038] The action command for the next time step and the core hidden state of the current time step are input into the reconstruction branch to generate the visual observation frame for the next time step. The specific process is as follows:

[0039] The motion command for the next time step is input into the second motion encoder to obtain the motion embedding for the next time step;

[0040] The action embedding of the next time step and the core hidden state of the current time step are input into the second cross-attention module. The action embedding of the next time step is used to obtain the key matrix and value matrix through linear projection. The core hidden state of the current time step is used to obtain the query matrix through linear projection. The action condition features of the next time step are calculated through the cross-attention mechanism.

[0041] The action condition features of the next time step are input into the conditional diffusion model as guiding conditions. Starting from randomly sampled Gaussian noise, iterative denoising is performed to generate a potential representation of the visual observation frame of the next time step.

[0042] The latent representation of the visual observation frame at the next time step is input into the second decoder to reconstruct the visual observation frame at the next time step.

[0043] Preferably, both the first decoder and the second decoder are decoders from variational autoencoders; both the first action encoder and the second action encoder are multilayer perceptron structures; and the conditional diffusion model uses the DiT architecture.

[0044] Secondly, the present invention provides a closed-loop strategy pre-action prediction device based on a world model, comprising:

[0045] The model building module is configured to build a policy model and a pre-trained world model. The policy model is trained based on the pre-trained world model using a closed-loop prediction process to obtain the trained policy model.

[0046] The prediction module is configured to acquire the visual observation frame at the first time step and perform an iterative prediction process for the first time step by combining the action command at the first time step with the zero-value vector and the core hidden state at the zeroth time step. During the iterative prediction process at the current time step, the pre-trained world model is used to obtain the core hidden state at the current time step based on the visual observation frame, action command, and core hidden state at the previous time step. The core hidden state at the current time step is input into the trained policy model to obtain the action command for the next time step. The pre-trained world model is then used to generate the visual observation frame for the next time step based on the action command for the next time step and the core hidden state at the current time step. The above iterative prediction process is repeated to obtain action commands for different time steps and send them to the machine to execute the corresponding actions.

[0047] Thirdly, the present invention provides an electronic device including one or more processors; and a storage device for storing one or more programs, wherein when the one or more programs are executed by the one or more processors, the one or more processors implement the method as described in any implementation of the first aspect.

[0048] Fourthly, the present invention provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the method as described in any of the implementations of the first aspect.

[0049] Fifthly, the present invention provides a computer program product, including a computer program that, when executed by a processor, implements the method as described in any of the implementations in the first aspect.

[0050] Compared with the prior art, the present invention has the following beneficial effects:

[0051] (1) The closed-loop policy forward action prediction method proposed in this invention injects visual embedding and action embedding into the prediction branch of the pre-trained world model, respectively, into the first adaptive normalization layer and the first cross attention module, which solves the problem of visual and action coordination inaccuracy in closed-loop policy forward prediction and reduces state prediction error.

[0052] (2) The closed-loop policy forward action prediction method based on the world model proposed in this invention uses the training branch of the pre-trained world model to restart the inference from the time step of the time step where the policy forward prediction failed, shortening the effective prediction path and suppressing the accumulation and amplification of long-term state error in the time-series inference stage.

[0053] (3) The closed-loop policy forward action prediction method proposed in this invention applies a mask constraint to the invalid frames of the time step after the successful time step during the training phase of the policy model, thereby reducing the interference of false success signals on gradient updates during policy prediction and improving the gradient stability and prediction reliability of the policy model. Attached Figure Description

[0054] To more clearly illustrate the technical solutions in the embodiments of the present invention, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying 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.

[0055] Figure 1 This is a flowchart illustrating the closed-loop policy foreshadowing action prediction method based on a world model, as an embodiment of this application.

[0056] Figure 2 This is a flowchart illustrating the prediction branch of the pre-trained world model in the closed-loop policy foreshadowing action prediction method based on the world model, as an embodiment of this application.

[0057] Figure 3 This is a flowchart illustrating the training branch of the pre-trained world model in the closed-loop policy foreshadowing action prediction method based on the world model, as described in an embodiment of this application.

[0058] Figure 4 This is a schematic diagram of the training process of the policy model in the closed-loop policy foreshadowing action prediction method based on the world model, which is an embodiment of this application.

[0059] Figure 5 This is a schematic diagram of a world model-based closed-loop strategy pre-action prediction device according to an embodiment of this application.

[0060] Figure 6 This is a schematic diagram of the hardware structure of an electronic device provided in an embodiment of the present invention. Detailed Implementation

[0061] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings. Obviously, the described embodiments are only a part of the embodiments of this invention, and not all of them. Based on the embodiments of this invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this invention.

[0062] Figure 1 An embodiment of this application illustrates a closed-loop policy foreshadowing action prediction method based on a world model, comprising the following steps:

[0063] S1. Construct a policy model and a pre-trained world model. Based on the pre-trained world model, use a closed-loop prediction process to train the policy model to obtain the trained policy model.

[0064] In a specific embodiment, the pre-trained world model includes a prediction branch, a training branch, and a reconstruction branch. The prediction branch includes a visual encoder, a first action encoder, a first adaptive normalization layer, and a first cross-attention module.

[0065] The visual observation frame and action command of the current time step, along with the core hidden state of the previous time step, are input into the prediction branch to obtain the core hidden state of the current time step. The specific process is as follows:

[0066] The visual observation frame and motion command of the current time step are input into the visual encoder and the first motion encoder, respectively, to obtain the visual embedding and motion embedding of the current time step.

[0067] The visual embedding of the current time step and the core hidden state of the previous time step are input into the first adaptive normalization layer. The visual embedding of the current time step is used as the input feature, and the core hidden state of the previous time step is used as the conditional information. The first intermediate feature of the current time step is modulated as shown in the following equation:

[0068] ;

[0069] Where t represents the current time step, and t-1 represents the previous time step. The visual embedding representing the current time step. This indicates the core hidden state of the previous time step. and represents the mean and standard deviation of the core hidden state at the previous time step, respectively; and These represent the scaling mapping function and the offset mapping function, respectively, both of which can be implemented by independent fully connected layers; and These represent the scaling vector and the offset vector, respectively. For element-wise multiplication operator, This represents the first intermediate feature at the current time step;

[0070] The action embedding and the first intermediate feature of the current time step are input into the first cross-attention module. The action embedding of the current time step is used to obtain the key matrix and value matrix through linear projection. The first intermediate feature of the current time step is used to obtain the query matrix through linear projection. The core hidden state of the current time step is calculated through the cross-attention mechanism.

[0071] Specifically, in the embodiments of this application, a real task trajectory dataset is first collected from the closed-loop policy pre-action prediction task. The real task trajectory dataset includes a sequence of real visual observation frames and a sequence of real action instructions from the start to the end of the task. The world model is then pre-trained using the real task trajectory dataset to obtain a pre-trained world model. The specific training process is as follows:

[0072] The first real visual observation frame in the sequence of real visual observation frames is used as the input to the world model. The current time step's visual observation frame is processed by the prediction branch to obtain the core hidden state of the current time step. The core hidden state of the current time step and the real action command of the next time step in the sequence of real action commands are processed by the reconstruction branch to obtain the next time step's visual observation frame. The real visual observation frame of the next time step is used as a supervision label and compared with the visual observation frame of the next time step predicted by the world model to construct a supervision loss function. The world model is pre-trained based on this supervision loss function to obtain a pre-trained world model.

[0073] The above process is not the focus of this application's embodiments. The focus of this application's embodiments is to train the policy model using a pre-trained world model to obtain a trained policy model. The structure of the policy model in this application's embodiments is not specifically limited; a multilayer perceptron structure or other reasonable structures can be used. During the training process of the policy model, the parameters to be optimized need to be adjusted. Meanwhile, the parameters of the pre-trained world model are frozen.

[0074] During the training of the policy model, the action instructions at the first time step and the core hidden state at time step zero Initialized using a zero vector, the visual observation frame at the first time step. Obtained by sampling from the real environment, the visual observation frame at the first time step. and action commands The input is fed into the pre-trained world model, which outputs the core hidden state of the first time step. The policy model takes the core hidden state of the first time step output by the pre-trained world model as input, outputs the action command of the next time step, and sends it to the machine for execution. The pre-trained world model takes the action command of the next time step as input and generates the visual observation frame of the next time step. The above steps are repeated to perform closed-loop prediction loop training on the policy model.

[0075] Specifically, the pre-trained world model includes three branches: prediction, training, and reconstruction. First, the prediction branch generates the core hidden state for the current time step based on the visual observation frame and action commands at that time step. Then, the core hidden state is combined with the training branch to construct the total loss function and its gradient during the policy model training process. The gradient of this total loss function is used to iteratively update the parameters of the policy model to minimize the total loss function. The reconstruction branch reconstructs the visual observation frame for the next time step using the action commands generated by the policy model.

[0076] refer to Figure 2 The prediction branch includes a visual encoder, a first action encoder, a first adaptive normalization layer, and a first cross-attention module. In one example, the visual encoding network employs... The first action encoding network uses a multilayer perceptron structure; other suitable structures can be chosen in other examples. The specific process is as follows:

[0077] Obtain the visual observation frame at the current time step. With action commands Each through a visual encoder With the first motion encoder Mapping to visual embedding With motion embedding embedding visuals Injected into the first adaptive normalization layer, utilizing visual embedding The core hidden state of the previous time step Adaptive normalization modulation is performed to obtain the first intermediate feature. Embedding actions and the first intermediate feature The input is fed into the first cross-attention module, where the core hidden state at the current time step is calculated using the cross-attention mechanism. In the embodiments of this application, visual embedding and action embedding are injected into the first adaptive normalization layer and the first cross-attention module, respectively, to enhance the visual-action co-consistency during the state evolution process.

[0078] In a specific embodiment, the training branch includes a state deviation measurement module, a keyframe retrieval module, and a state reset module; the state deviation measurement module includes a second adaptive normalization layer, a multilayer perceptron, a first decoder, and a linear projection layer.

[0079] The training process of the policy model includes:

[0080] First, execute the closed-loop prediction process. The steps of the closed-loop prediction process are as follows:

[0081] The visual observation frame and action command of the current time step in the training sample, as well as the core hidden state of the previous time step, are passed through the prediction branch to obtain the core hidden state of the current time step.

[0082] The core hidden state of the previous time step, the core hidden state of the current time step, and the action command from the training samples are input into the training branch. First, they pass through the state bias measurement module. Then, the core hidden state of the previous time step and the action command from the current time step are input into the second adaptive normalization layer. The core hidden state of the previous time step is used as the input feature, and the action command from the current time step is used as the conditional information to modulate and obtain the second intermediate feature of the current time step. The second intermediate feature of the current time step is then input into the multilayer perceptron to map and obtain the reference state of the current time step.

[0083] The core hidden state at the current time step is passed sequentially through the first decoder and the linear projection layer to map the predicted state at the current time step.

[0084] Calculate the deviation between the predicted state and the reference state at the current time step; in response to the determination that the deviation values ​​of N consecutive time steps t-N+1, t-N+2, ..., t exceed the deviation threshold, trigger the state adjustment process. The current time step is used as the time step in which the strategy's forward prediction fails, and the reference state of the current time step is used as the reference state of the time step in which the strategy's forward prediction fails.

[0085] During the state adjustment process The time step where the strategy-predicted failure occurred and its reference state are input into the keyframe retrieval module. The prediction states of the k time steps preceding the time step where the strategy-predicted failure occurred are used to form a neighboring candidate window set for the time step where the strategy-predicted failure occurred. The prediction state of the optimal keyframe is retrieved from this neighboring candidate window set, as shown in the following formula:

[0086] ;

[0087] in, This represents the set of neighboring candidate windows at the time step where the strategy's prediction failed. The index i represents the prediction state in the set of neighboring candidate windows at the time step where the strategy prediction failed, where i = tk, t-k+1, ..., t-1; This represents the prediction state of the optimal keyframe. The independent variable represents the value at which the function reaches its minimum. This indicates the reference state at the time step in which the strategy's prediction failed. This is the timing penalty coefficient. For window length, Represents the L2 norm;

[0088] The prediction state of the optimal keyframe is input into the state reset module, the core hidden state of the time step corresponding to the prediction state of the optimal keyframe is extracted, and it is forcibly overwritten on the core hidden state of the current time step of the time step that failed as the strategy forward prediction.

[0089] The closed-loop prediction process is repeated, and the core hidden state of the current time step, which is the time step where the strategy foreshadowing prediction fails, is adjusted to obtain the core hidden state of each time step.

[0090] The success probability of the task at each time step is obtained by mapping the core hidden state at each time step, as shown in the following formula:

[0091] ;

[0092] in, and These represent the weight matrix and the bias term, respectively. This indicates the probability of task success at the current time step;

[0093] The time step with a success probability greater than the success determination threshold is set as the success time step. The visual observation frames of the success time step and the time steps before it are taken as valid frames and the corresponding mask is set to 1. The visual observation frames of the time steps after the success time step are taken as invalid frames and the corresponding mask is set to 0.

[0094] The total loss function and its gradient used during the training of the strategy model are calculated using the following formula:

[0095] ;

[0096] ;

[0097] ;

[0098] in, Represents the total loss function. This represents the cumulative discounted reward from the successful time step (l) to the successful time step T, where j represents the time step variable, j = l, l+1, ..., T. j-l Represents time step variables j With one of the time steps l difference, Indicates the discount factor. A mask representing the successful time step and one of the time steps l preceding it. This represents the core hidden state of time step l, which is one of the time steps preceding the successful time step. After being input into the policy model, the action command for the next time step is output. The conditional probability, Indicates the parameter to be optimized is The strategy model; This represents the probability of task success for time step variable j; The parameters to be optimized in the strategy model Gradient operator for variables;

[0099] The parameters to be optimized in the policy model are updated using the gradient of the total loss function, resulting in a trained policy model.

[0100] In a specific embodiment, both the first adaptive normalization layer and the second adaptive normalization layer are adaptive normalization layers; the visual encoder adopts the Vision Transformer structure, and the deviation value is the square of the Euclidean distance.

[0101] For details, please refer to Figure 3 The training branches of the pre-trained world model include a state bias measurement module, a keyframe retrieval module, and a state reset module. The state bias measurement module includes a second adaptive normalization layer, a multilayer perceptron, a first decoder, and a linear projection layer. Both the first and second adaptive normalization layers employ AdaLN (Adaptive Layer Normalization), and their specific calculation process is described in the section on the first adaptive normalization layer. The closed-loop prediction process is then repeatedly executed, combining the training branches of the pre-trained world model and the policy model, to train the policy model. The steps of the closed-loop prediction process are as follows:

[0102] First, the prediction branch of the pre-trained world model is used to predict the core hidden state of the current time step based on the visual observation frame and action command of the current time step in the training samples and the core hidden state of the previous time step.

[0103] In the state deviation measurement module, the action command of the current time step is first processed using the second adaptive normalization layer. Encoded as modulation parameters, representing the core hidden state at the current time step. Feature modulation is performed to obtain the second intermediate feature at the current time step; this second intermediate feature at the current time step is then mapped to the reference state at the current time step using a multilayer perceptron. The reference state at the current time step It serves as a short-term reliable benchmark for comparing deviations from the predicted state at the current time step.

[0104] Then hide the core state at the current time step. The data is sequentially passed through the first decoder and the linear projection layer, and projected onto the observation space to extract the predicted state at the current time step. And quantify its reference state relative to the current time step. The degree of deviation is determined, and the deviation value is obtained as shown in the following formula:

[0105] ;

[0106] in, This represents the predicted state at the current time step and the reference state at the current time step. The deviation value.

[0107] If the deviation value is continuous The deviation value of the time step exceeds the deviation threshold. Then the state adjustment process is triggered. At this moment, the current time step The time step where the strategy predicts failure, and the corresponding prediction state. The reference state is the predicted state when the strategy fails to predict the outcome. This serves as a reference state for predicting the failure of the strategy.

[0108] During the state adjustment process Define the current time step as the time step for predicting the failure of the strategy. Based on this, construct the corresponding set of neighboring candidate windows. ,in For a preset window length, the keyframe retrieval module retrieves the prediction state of the optimal keyframe within the local search space. During the retrieval process, the reference state of relocation anchoring is used as the screening criterion to quantify the fit between the candidate frames in the neighboring candidate window set and the model inference logic, and the time deviation is controlled by the time penalty coefficient.

[0109] Extract the optimal keyframe prediction state in the state reset module. The core hidden state at the corresponding time step Path correction is achieved by forcibly overwriting the core hidden state of the current time step where the policy foreshadowing prediction failed. After the state is reset, a new round of closed-loop prediction is restarted. Through this closed-loop prediction process, the core hidden states in the pre-trained world model that failed the policy foreshadowing prediction are adjusted to the core hidden states that succeeded. Policy foreshadowing is performed from visual observation frames near the time step where the policy foreshadowing prediction failed, shortening the effective prediction path, achieving dynamic repair of temporal consistency, and avoiding long-term error accumulation.

[0110] refer to Figure 4 Furthermore, the core hidden state mapping of the current time step is adjusted using the training branch of the pre-trained world model to obtain the task success probability of the current time step. If the task success probability of the current time step is greater than the success determination threshold, then the current time step is a successful time step. Visual observation frames at and before the successful time step are considered valid frames, while those after are considered invalid frames. Binary masks are applied to both valid and invalid frames to mark the validity of the observation frame sequence. If a frame is determined to be valid, the corresponding mask value is 1; if it is determined to be invalid, the corresponding mask value is 0. The gradient flow for the corresponding visual observation frame is then enabled or disabled using this mask. By imposing masking constraints on invalid frames, the interference of false success signals on the gradient updates of the policy model is reduced.

[0111] In strategy model Parameters to be optimized During the update phase, the cumulative discounted reward from one of the time steps l (including the successful time step and the time steps before it) to the successful time step T is first calculated, so that the gradient direction is biased towards the action command that can obtain a higher reward. Then, this mask is used to weight and correct the gradient of the total loss function to ensure that the gradient of invalid frames does not participate in the parameters to be optimized. The iterative process, through repeated updates, yields the optimized parameters of the policy model. Because a masking constraint is introduced during the training phase of the policy model to mask the gradients of invalid frames, the parameters to be optimized... It is formed solely by the accumulation of gradient signals from valid frames, thus making the action command output at key decision points more deterministic during the inference phase.

[0112] The policy model generates the action command for the next time step based on the core hidden state of the current time step. Then, the reconstruction branch of the pre-trained world model generates the visual observation frame for the next time step based on the action command for the next time step and the core hidden state of the current time step. This process is repeated to train the policy model and optimize its parameters. The optimization process is repeated to obtain the trained policy model. In one example, the optimizer used during the training of the policy model employs the Adam algorithm with a weight decay rate of 5 × 10⁻⁶. −4 The momentum is set to 0.9, and the batch size is 64.

[0113] The training branches of the pre-trained world model are only used during the training phase of the policy model and not during the inference phase. During the inference phase, only the prediction and reconstruction branches of the pre-trained world model are used, while all three branches of the pre-trained world model are used during the training phase of the policy model.

[0114] S2: Obtain the visual observation frame of the first time step, and combine the action command of the first time step with the zero vector value and the core hidden state of the zeroth time step to perform the iterative prediction process of the first time step; In the iterative prediction process of the current time step, use the pre-trained world model to obtain the core hidden state of the current time step based on the visual observation frame of the current time step, the action command, and the core hidden state of the previous time step; Input the core hidden state of the current time step into the trained policy model to obtain the action command of the next time step, and use the pre-trained world model to generate the visual observation frame of the next time step based on the action command of the next time step and the core hidden state of the current time step; Repeat the above iterative prediction process to obtain the action commands of different time steps and send them to the machine to execute the corresponding actions.

[0115] In a specific embodiment, the reconstruction branch includes a second action encoder, a second cross-attention module, a conditional diffusion model, and a second decoder;

[0116] The action command for the next time step and the core hidden state of the current time step are input into the reconstruction branch to generate the visual observation frame for the next time step. The specific process is as follows:

[0117] The motion command for the next time step is input into the second motion encoder to obtain the motion embedding for the next time step;

[0118] The action embedding of the next time step and the core hidden state of the current time step are input into the second cross-attention module. The action embedding of the next time step is used to obtain the key matrix and value matrix through linear projection. The core hidden state of the current time step is used to obtain the query matrix through linear projection. The action condition features of the next time step are calculated through the cross-attention mechanism.

[0119] The action condition features of the next time step are input into the conditional diffusion model as guiding conditions. Starting from randomly sampled Gaussian noise, iterative denoising is performed to generate a potential representation of the visual observation frame of the next time step.

[0120] The latent representation of the visual observation frame at the next time step is input into the second decoder to reconstruct the visual observation frame at the next time step.

[0121] In a specific embodiment, both the first decoder and the second decoder are decoders in a variational autoencoder; both the first action encoder and the second action encoder are multilayer perceptron structures; and the conditional diffusion model adopts a DiT architecture.

[0122] Specifically, the trained policy model is deployed. During the inference phase, the visual observation frame of the first time step to be predicted is acquired and combined with the action command of the first time step and the core hidden state of the zeroth time step to execute the iterative prediction process of the first time step. In the iterative prediction process of the current time step, the prediction branch of the pre-trained world model is used to predict the core hidden state of the current time step based on the visual observation frame, action command, and core hidden state of the previous time step. The core hidden state of the current time step is input into the trained policy model, which can accurately predict the action command of the next time step. The reconstruction branch of the pre-trained world model is then used to generate the visual observation frame of the next time step based on the action command of the next time step and the core hidden state of the current time step. The visual observation frame and action command of the next time step are used as the visual observation frame and action command of the current time step, and the core hidden state of the current time step is used as the core hidden state of the previous time step. By repeating the above iterative prediction process, the action command of each time step can be predicted and sent to the machine to execute the corresponding action. Since the policy model has learned to generate action instructions with a high success rate at key decision points during the training phase, it can be directly controlled based on the stable core hidden state provided by the pre-trained world model during inference.

[0123] This invention addresses the problem of predicted states deviating from the actual physical states in existing world models. Specifically, during the feature generation stage, visual embeddings and action embeddings are injected into the first adaptive normalization layer and the first cross-attention module, respectively, to enhance visual-action co-consistency during state evolution and reduce single-step prediction errors. Furthermore, a policy pre-expansion method based on keyframe relocalization is proposed. When the predicted state meets the conditions for policy pre-expansion prediction failure, the inference is restarted from the neighboring keyframe of the failed time step, shortening the effective prediction path and suppressing the accumulation and amplification of errors in long-term inference. Finally, this invention applies a gradient mask to invalid frames after successful time steps, reducing the interference of false success signals on gradient updates during policy prediction, forming a multi-stage illusion suppression mechanism of "feature alignment—path correction—gradient constraint".

[0124] Further reference Figure 5As an implementation of the methods shown in the above figures, this application provides an embodiment of a closed-loop policy foreshadowing action prediction device based on a world model. This device embodiment is similar to... Figure 1 Corresponding to the method embodiments shown, this device can be specifically applied to various electronic devices.

[0125] This application provides a closed-loop policy foreshadowing action prediction device based on a world model, including:

[0126] Model building module 1 is configured to build a policy model and a pre-trained world model. The policy model is trained based on the pre-trained world model using a closed-loop prediction process to obtain the trained policy model.

[0127] Prediction module 2 is configured to acquire the visual observation frame of the first time step and perform an iterative prediction process for the first time step by combining the action command of the first time step with the zero vector value and the core hidden state of the zeroth time step. During the iterative prediction process of the current time step, the pre-trained world model is used to obtain the core hidden state of the current time step based on the visual observation frame, action command, and core hidden state of the previous time step. The core hidden state of the current time step is input into the trained policy model to obtain the action command of the next time step. The pre-trained world model is used to generate the visual observation frame of the next time step based on the action command of the next time step and the core hidden state of the current time step. The above iterative prediction process is repeated to obtain the action commands of different time steps and send them to the machine to execute the corresponding actions.

[0128] Figure 6 This is a schematic diagram of the hardware structure of an electronic device provided in an embodiment of the present invention. For example... Figure 6 As shown, the electronic device of this embodiment includes a processor 601 and a memory 602; wherein the memory 602 is used to store computer execution instructions; and the processor 601 is used to execute the computer execution instructions stored in the memory to implement the various steps performed by the electronic device in the above embodiment. For details, please refer to the relevant descriptions in the foregoing method embodiments.

[0129] Alternatively, the memory 602 can be either standalone or integrated with the processor 601.

[0130] When the memory 602 is set up independently, the electronic device also includes a bus 603 for connecting the memory 602 and the processor 601.

[0131] This invention also provides a computer storage medium storing computer execution instructions, which, when executed by processor 601, implement the above method.

[0132] This invention also provides a computer program product, including a computer program that, when executed by a processor 601, implements the above-described method.

[0133] In the embodiments provided by this invention, it should be understood that the disclosed devices and methods can be implemented in other ways. For example, the device embodiments described above are merely illustrative; for instance, the division of modules is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple modules may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be indirect coupling or communication connection through some interfaces, devices, or modules, and may be electrical, mechanical, or other forms.

[0134] The modules described as separate components may or may not be physically separate. The components shown as modules may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to implement the solution of this embodiment according to actual needs.

[0135] Furthermore, the functional modules in the various embodiments of this invention can be integrated into one processing unit, or each module can exist physically separately, or two or more modules can be integrated into one unit. The unit formed by the above modules can be implemented in hardware or in the form of hardware plus software functional units.

[0136] The integrated modules described above, implemented as software functional modules, can be stored in a computer-readable storage medium. These software functional modules, stored in a storage medium, include several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) or processor 601 to execute some steps of the methods of the various embodiments of this application.

[0137] It should be understood that the processor 601 described above can be a Central Processing Unit (CPU), or other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), etc. The general-purpose processor can be a microprocessor, or the processor 601 can be any conventional processor 601. The steps of the method disclosed in this invention can be directly manifested as the hardware processor 601 executing the steps, or as a combination of hardware and software modules within the processor 601 executing the steps.

[0138] The memory 602 may include high-speed RAM memory, and may also include non-volatile memory NVM, such as at least one disk storage device, and may also be a USB flash drive, portable hard drive, read-only memory, disk or optical disc, etc.

[0139] Bus 603 can be an Industry Standard Architecture (ISA), a Peripheral Component Interconnect (PCI) bus, or an Extended Industry Standard Architecture (EISA) bus, etc. Bus 603 can be divided into address bus, data bus, control bus, etc. For ease of illustration, the bus 603 in the accompanying drawings of this application is not limited to only one bus 603 or one type of bus 603.

[0140] The aforementioned storage medium can be implemented from any type of volatile or non-volatile storage device or a combination thereof, such as static random access memory (SRAM), electrically erasable programmable read-only memory (EEPROM), erasable programmable read-only memory (EPROM), programmable read-only memory (PROM), read-only memory (ROM), magnetic storage, flash memory, magnetic disk, or optical disk. The storage medium can be any available medium accessible to general-purpose or special-purpose computers.

[0141] An exemplary storage medium is coupled to a processor 601, enabling the processor 601 to read information from and write information to the storage medium. Alternatively, the storage medium can be an integral part of the processor 601. The processor 601 and the storage medium can reside in an application-specific integrated circuit (ASIC). Alternatively, the processor 601 and the storage medium can exist as discrete components in an electronic device or a host device.

[0142] Those skilled in the art will understand that all or part of the steps of the above-described method embodiments can be implemented by hardware related to program instructions. The aforementioned program can be stored in a computer-readable storage medium. When executed, the program performs the steps of the above-described method embodiments; and the aforementioned storage medium includes various media capable of storing program code, such as ROM, RAM, magnetic disks, or optical disks.

[0143] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features therein. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. A closed-loop policy foreshadowing action prediction method based on a world model, characterized in that, Includes the following steps: Construct a policy model and a pre-trained world model, and train the policy model using a closed-loop prediction loop process based on the pre-trained world model to obtain a trained policy model. The system acquires the visual observation frame at the first time step and performs an iterative prediction process for the first time step by combining the action command at the first time step with a value of zero and the core hidden state at the zeroth time step. During the iterative prediction process at the current time step, the pre-trained world model is used to obtain the core hidden state of the current time step based on the visual observation frame, action command, and core hidden state of the previous time step. The core hidden state of the current time step is input into the trained policy model to obtain the action command for the next time step. The pre-trained world model is then used to generate the visual observation frame for the next time step based on the action command for the next time step and the core hidden state of the current time step. This iterative prediction process is repeated to obtain action commands for different time steps, which are then sent to the machine to execute the corresponding actions.

2. The closed-loop policy foreshadowing action prediction method based on a world model according to claim 1, characterized in that, The pre-trained world model includes a prediction branch, a training branch, and a reconstruction branch. The prediction branch includes a visual encoder, a first action encoder, a first adaptive normalization layer, and a first cross-attention module. The visual observation frame and action command of the current time step, as well as the core hidden state of the previous time step, are input into the prediction branch to obtain the core hidden state of the current time step. The specific process is as follows: The visual observation frame and motion command of the current time step are input into the visual encoder and the first motion encoder, respectively, to obtain the visual embedding and motion embedding of the current time step. The visual embedding of the current time step and the core hidden state of the previous time step are input into the first adaptive normalization layer. The visual embedding of the current time step is used as the input feature, and the core hidden state of the previous time step is used as the conditional information. The first intermediate feature of the current time step is modulated as shown in the following equation: ； Where t represents the current time step, and t-1 represents the previous time step. The visual embedding representing the current time step. This indicates the core hidden state of the previous time step. and represents the mean and standard deviation of the core hidden state at the previous time step, respectively; and These represent the scaling mapping function and the offset mapping function, respectively, both of which can be implemented by independent fully connected layers; and These represent the scaling vector and the offset vector, respectively. For element-wise multiplication operator, This represents the first intermediate feature at the current time step; The action embedding and the first intermediate feature of the current time step are input into the first cross-attention module. The action embedding of the current time step is used to obtain the key matrix and value matrix through linear projection. The first intermediate feature of the current time step is used to obtain the query matrix through linear projection. The core hidden state of the current time step is calculated through the cross-attention mechanism.

3. The closed-loop policy foreshadowing action prediction method based on a world model according to claim 2, characterized in that, The training branch includes a state deviation measurement module, a keyframe retrieval module, and a state reset module; The state deviation measurement module includes a second adaptive normalization layer, a multilayer perceptron, a first decoder, and a linear projection layer. The training process of the policy model includes: First, a closed-loop prediction loop process is executed. The steps of the closed-loop prediction loop process are as follows: The visual observation frame and action command of the current time step in the training sample, as well as the core hidden state of the previous time step, are passed through the prediction branch to obtain the core hidden state of the current time step. The core hidden state of the previous time step, the core hidden state of the current time step, and the action command from the training samples are input into the training branch. First, the training branch passes through the state deviation measurement module. Then, the core hidden state of the previous time step and the action command from the current time step are input into the second adaptive normalization layer. The core hidden state of the previous time step is used as the input feature, and the action command from the current time step is used as the conditional information. The second intermediate feature of the current time step is modulated to obtain the second intermediate feature of the current time step. The second intermediate feature of the current time step is input into the multilayer perceptron to map the reference state of the current time step. The core hidden state of the current time step is sequentially passed through the first decoder and the linear projection layer to map the predicted state of the current time step. Calculate the deviation between the predicted state and the reference state at the current time step; in response to the determination that the deviation values ​​of N consecutive time steps t-N+1, t-N+2, ..., t exceed the deviation threshold, trigger the state adjustment process. The current time step is used as the time step in which the strategy's forward prediction fails, and the reference state of the current time step is used as the reference state of the time step in which the strategy's forward prediction fails. During the state adjustment process, the time step where the strategy prediction failed and its reference state are input into the keyframe retrieval module. The prediction states of the k time steps before the time step where the strategy prediction failed are used to form a neighboring candidate window set for the time step where the strategy prediction failed. The prediction state of the optimal keyframe is retrieved within the neighboring candidate window set for the time step where the strategy prediction failed, as shown in the following formula: ； in, This represents the set of neighboring candidate windows at the time step where the strategy's prediction failed. The index i represents the prediction state in the set of neighboring candidate windows at the time step where the strategy prediction failed, where i = tk, t-k+1, ..., t-1; This represents the prediction state of the optimal keyframe. The independent variable represents the value at which the function reaches its minimum. This indicates the reference state at the time step in which the strategy's prediction failed. This is the timing penalty coefficient. For window length, Represents the L2 norm; The prediction state of the optimal keyframe is input into the state reset module, the core hidden state of the time step corresponding to the prediction state of the optimal keyframe is extracted, and it is forcibly overwritten on the core hidden state of the current time step of the time step that is the time step of the strategy forward prediction failure. Repeat the closed-loop prediction process, and adjust the core hidden state of the current time step, which is the time step where the strategy foreshadowing prediction fails, to obtain the core hidden state of each time step. The success probability of the task at each time step is obtained by mapping the core hidden state at each time step, as shown in the following formula: ； in, and These represent the weight matrix and the bias term, respectively. This indicates the probability of task success at the current time step; The time step in which the success probability of the task is greater than the success determination threshold is set as the success time step. The visual observation frames of the success time step and the time steps before it are taken as valid frames and the corresponding mask is set to 1. The visual observation frames of the time steps after the success time step are taken as invalid frames and the corresponding mask is set to 0. The total loss function and its gradient used during the training of the policy model are calculated using the following formula: ； ； ； in, Represents the total loss function. This represents the cumulative discounted reward from one of the time steps l (including the successful time step and the time steps preceding it) to the successful time step T, where j represents the time step variable, j = l, l+1, ..., T. j-l Represents time step variables j With one of the time steps l difference, Indicates the discount factor. The mask representing the successful time step and one of the time steps preceding it, l. This represents the core hidden state of one of the time steps l, including the successful time step and the time steps preceding it. After being input into the policy model, the action command for the next time step is output. The conditional probability, Indicates the parameter to be optimized is The strategy model; This represents the probability of task success for time step variable j; The parameters to be optimized in the strategy model Gradient operator for variables; The parameters to be optimized in the policy model are updated using the gradient of the total loss function to obtain the trained policy model.

4. The closed-loop policy foreshadowing action prediction method based on a world model according to claim 3, characterized in that, Both the first and second adaptive normalization layers are adaptive normalization layers; the visual encoder adopts a Vision Transformer structure, and the deviation value is the square of the Euclidean distance.

5. The closed-loop policy foreshadowing action prediction method based on a world model according to claim 3, characterized in that, The reconstruction branch includes a second action encoder, a second cross-attention module, a conditional diffusion model, and a second decoder; The action command for the next time step and the core hidden state of the current time step are input into the reconstruction branch to generate the visual observation frame for the next time step. The specific process is as follows: The motion command for the next time step is input into the second motion encoder to obtain the motion embedding for the next time step; The action embedding of the next time step and the core hidden state of the current time step are input into the second cross-attention module. The action embedding of the next time step is used to obtain the key matrix and value matrix through linear projection. The core hidden state of the current time step is used to obtain the query matrix through linear projection. The action condition features of the next time step are calculated through the cross-attention mechanism. The action condition features of the next time step are input into the conditional diffusion model as guiding conditions, and iterative denoising is performed starting from randomly sampled Gaussian noise to generate a potential representation of the visual observation frame of the next time step. The latent representation of the visual observation frame at the next time step is input into the second decoder to reconstruct the visual observation frame at the next time step.

6. The closed-loop policy foreshadowing action prediction method based on a world model according to claim 5, characterized in that, Both the first decoder and the second decoder are decoders in variational autoencoders; both the first action encoder and the second action encoder are multilayer perceptron structures; the conditional diffusion model adopts the DiT architecture.

7. A closed-loop strategy pre-action prediction device based on a world model, characterized in that, include: The model building module is configured to build a policy model and a pre-trained world model, and to train the policy model based on the pre-trained world model using a closed-loop prediction loop process to obtain a trained policy model. The prediction module is configured to acquire the visual observation frame of the first time step and perform an iterative prediction process for the first time step by combining the action command of the first time step with the zero-value vector and the core hidden state of the zeroth time step. During the iterative prediction process of the current time step, the pre-trained world model is used to obtain the core hidden state of the current time step based on the visual observation frame, action command, and core hidden state of the previous time step. The core hidden state of the current time step is input into the trained policy model to obtain the action command for the next time step. The pre-trained world model is used to generate the visual observation frame for the next time step based on the action command for the next time step and the core hidden state of the current time step. The above iterative prediction process is repeated to obtain action commands for different time steps and send them to the machine to execute the corresponding actions.

8. An electronic device, comprising: One or more processors; A storage device for storing one or more programs, characterized in that, when the one or more programs are executed by the one or more processors, the one or more processors implement the method as described in any one of claims 1-6.

9. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the method as described in any one of claims 1-6.

10. A computer program product, comprising a computer program, characterized in that, When the computer program is executed by a processor, it implements the method as described in any one of claims 1-6.

Images