An unsupervised reinforcement learning method and system based on world model extension

Abstract

The application provides a kind of based on world model extension's unsupervised reinforcement learning method and system, belong to computer deep learning field.The method obtains the history offline data set D of expert strategy and environment interaction after being trained;World model neural network is constructed, and based on D is trained;Unsupervised reinforcement learning model is constructed again, and initial current predicted state is set;Current predicted state, skill vector z are input into policy network to obtain action;Current predicted state and action are input into mature world model neural network, and next predicted state is obtained;After obtaining complete trajectory data, it is stored in replay buffer;Part of trajectory data is randomly sampled, and unsupervised reinforcement learning model is trained;Current environment state and the skill vector z of agent are input into mature unsupervised reinforcement learning model, output action parameter, and control agent to execute corresponding action.The application improves the generalization ability and learning efficiency of unsupervised reinforcement learning.

Description

Technical Field

[0001] This invention belongs to the field of computer deep learning and relates to an unsupervised reinforcement learning method and system based on world model extension. Background Technology

[0002] Reinforcement learning, as a method for learning optimal decisions through interaction with the environment, has demonstrated great potential in many fields. However, traditional reinforcement learning algorithms typically require large amounts of online interaction data, which often leads to low sample efficiency in practical applications. Especially in complex environments, agents need to spend a significant amount of time exploring to learn effective policies, greatly limiting the scope of practical applications of reinforcement learning. To overcome this limitation, researchers have recently focused on how to utilize offline data for pre-training, thereby significantly improving the sample efficiency of reinforcement learning. Against this backdrop, unsupervised reinforcement learning has gradually become a research hotspot. Its core idea is to learn general skills or representations through offline data without relying on external reward signals, providing a good initialization for subsequent online tasks.

[0003] The key advantage of unsupervised reinforcement learning lies in its ability to autonomously discover useful behavioral patterns from historical data without requiring manually designed reward functions. This approach is particularly suitable for real-world task scenarios where reward signals are sparse or difficult to define. Through unsupervised pre-training, agents can learn a series of basic skills and intrinsic representations of environmental dynamics. This prior knowledge can significantly accelerate the learning process for subsequent specific tasks. Skill learning, as an important branch of unsupervised reinforcement learning, aims to extract reusable basic behavioral units from historical experience. These skills can be flexibly combined like "building blocks" to provide efficient solutions for subsequent tasks. For example, in the field of robot control, agents can learn various basic action skills in an unsupervised manner, and in actual tasks, they can quickly adapt to new requirements simply by combining these skills. Furthermore, unsupervised learning can also discover solution strategies that humans may not have thought of, providing new approaches to complex problems.

[0004] In existing technologies, researchers have proposed numerous methods for learning skills through unsupervised reinforcement learning on offline data. Some scholars have pioneered the introduction of skill learning into offline settings using VAEs to efficiently capture behavioral skills from offline data, theoretically proving its superiority over direct learning in the action space. Others, addressing the limited representational limitations of previous methods in recovering the original policy space, have proposed using a flow model to learn an invertible function, mapping latent embeddings to temporally extended actions. Still others have proposed obtaining behavioral priors by learning representations of input-output relationships, enabling faster learning of new tasks without hindering the agent's ability to learn new behaviors. To address the underdeterminacy of previous methods that maximize likelihood for skill learning, researchers have introduced an additional objective, biasing skill learning towards acquiring skills that accelerate new task learning. Furthermore, some researchers have proposed mapping states to a Hilbert space, learning structured representations that preserve environmental structure to capture diverse, optimal, and long-term behaviors from unlabeled offline data. Finally, to enable agents to learn diverse and useful behaviors from unsupervised offline data and perform rapid online learning, researchers have proposed imbuing the dataset with random intentions to encourage the diversity and usefulness of extracted skills.

[0005] However, the skill learning capabilities of these methods are severely limited by the coverage of the training data. When faced with new scenarios that exceed the original data distribution, existing skills often struggle to generalize effectively. Secondly, most methods focus only on the quality of skill reconstruction itself, neglecting the dynamic adaptation relationship between skills and the environment, resulting in unstable performance of learned skills in actual interactions. More importantly, existing methods typically treat skill learning as a static process, failing to establish a mechanism for the dynamic expansion of the skill set as the environment changes, thus leading to poor application performance in complex and open environments. Summary of the Invention

[0006] In view of the above-mentioned defects or deficiencies in the prior art, the present invention aims to provide an unsupervised reinforcement learning method and system based on world model extension. It achieves comprehensive modeling of environmental dynamics through world model integration architecture, ensures high controllability of learned skills by using explicit alignment mechanism between state changes and skill direction, and adopts adaptive constraint strategy to balance exploration efficiency and stability, so as to improve the generalization ability and learning efficiency of skills.

[0007] To achieve the above objectives, the embodiments of the present invention adopt the following technical solutions:

[0008] In a first aspect, embodiments of the present invention provide an unsupervised reinforcement learning method based on an extended world model, comprising the following steps:

[0009] Step S1: Obtain the historical offline dataset D of the trained expert policy and its interaction with the environment;

[0010] Step S2: Construct the world model neural network;

[0011] Step S3: Based on the historical offline dataset D, train the world model neural network to obtain a mature world model neural network;

[0012] Step S4: Construct an unsupervised reinforcement learning model and set the initial current prediction state;

[0013] Step S5: Input the current predicted state and skill vector z into the policy network of the unsupervised reinforcement learning model to obtain the action;

[0014] Step S6: Input the current predicted state and action into the mature world model neural network to obtain the next predicted state; and determine whether the maximum step size or the termination state has been reached; if not, take the next predicted state as the current predicted state and return to step S5; if yes, obtain the complete trajectory data, store it in the playback buffer, and proceed to step S7.

[0015] Step S7: Randomly sample a portion of trajectory data from the playback buffer to train the unsupervised reinforcement learning model and obtain a mature unsupervised reinforcement learning model.

[0016] Step S8: Input the current environment state and the agent's skill vector z into the mature unsupervised reinforcement learning model, output the agent's action parameters, and control the agent to perform corresponding actions according to the action parameters.

[0017] In a preferred embodiment of the present invention, step S2, which involves constructing a world model neural network, includes:

[0018] An integrated model architecture is adopted to construct a world model neural network; the world model neural network contains n sub-models with identical structures, and each sub-model consists of a main network and a prior network, with the current environment state s as the input to both networks. t and the action a performed by the intelligent agent t The output is the predicted difference Δs between the predicted next state and the current state. t Furthermore, the parameters of the prior network for each sub-model are fixed but different; among them, the current state and actions d and m represent the dimensions of the state and action, respectively.

[0019] In a preferred embodiment of the present invention, the sub-model adopts a multilayer perceptron structure, and the multilayer perceptron structure includes three hidden layers; the integrated model architecture adopts a parallel computing architecture, and the constructed world model neural network supports multiple sub-models to perform state prediction simultaneously.

[0020] As a preferred embodiment of the present invention, step S2 trains the world model neural network based on the historical offline dataset D, which specifically includes:

[0021] According to the environmental states at different moments and the actions executed by the agent in the historical offline dataset D, calculate the difference between the next state and the current state at each moment, incorporate the differences into the dataset to construct a new dataset D*, divide D* into a training set, a validation set and a test set; then train the main network, which is a sub-model of the world model neural network, based on the training set in D*, and perform validation and testing through the validation set and the test set respectively.

[0022] As a preferred embodiment of the present invention, the training of the main network, which is a sub-model of the world model neural network, based on the training set in D* adopts a differentiated prior method to train the world model neural network;

[0023] During training, the prior network of each sub-model does not participate in direct training, its parameters are fixed and different, but the predicted difference of the prior network is obtained by prediction according to the input value; the main network of each sub-model directly participates in training, after updating the parameters of the main network in each training round, the predicted difference of the main network is obtained by prediction according to the input value; the output of the sub-model is the weighted sum of the predicted differences output by the main network and the prior network;

[0024] During training, the composite loss function adopted comprises four components: a weighted mean square error loss term for minimizing state prediction deviation, a logarithmic variance loss term for quantifying prediction uncertainty, a model parameter regularization term for preventing overfitting, and a logarithmic variance range constraint term for ensuring numerical stability.

[0025] As a preferred embodiment of the present invention, the unsupervised reinforcement learning model constructed in step S4 comprises at least three components: a state representation network a policy network π and a Q-function network; wherein,

[0026] the state representation network adopts a multi-layer perceptron (MLP) structure, comprising two hidden layers, and its input is an original state and its output is a low-dimensional representation where k<<d, the network uses a ReLU activation function to ensure training stability; the state representation network is used to extract key features of the state;

[0027] the policy network π adopts an MLP structure, comprising two hidden layers, and its input is a state representation and a skill vector z, and outputs the mean and variance parameters of the action distribution; to constrain the action range, a tanh activation function is used in the output layer, and a reparameterization trick is adopted to implement a differentiable stochastic policy;

[0028] The Q-function network uses an MLP structure and contains two hidden layers, representing the input state. Skill vector z and action a, output state-skill value estimate.

[0029] In a preferred embodiment of the present invention, step S6 inputs the current predicted state and action into a mature world model neural network to obtain the next predicted state. Specifically, this includes: sampling a skill vector z from a preset Gaussian distribution as a strategy guidance signal for the current round; interacting with the world model neural network based on the strategy network π(a|s, z), whereby the strategy network generates action a according to the current predicted state and skill vector z, and the world model neural network receives action a and the current predicted state, outputting the prediction difference Δs. t The current predicted state is added to the difference between the predicted state and the model output to obtain the next predicted state; finally, complete trajectory data τ={…,(s t a t s t+1 ), ...} and store them in the playback buffer pool.

[0030] In a preferred embodiment of the present invention, step S7 trains the unsupervised reinforcement learning model by calculating the loss value of the unsupervised reinforcement learning model based on the data in the replay buffer and updating the parameters; it determines whether the maximum number of training rounds has been reached; if not, it returns to step S5; if it has been reached, the training is completed and a mature unsupervised reinforcement learning model is obtained.

[0031] In a preferred embodiment of the present invention, a state representation network is included during parameter updates. The parameters of the three components—the policy network (π), the Q-function network, and the parameter update process—are updated as follows:

[0032] First, a loss function is learned using state representation. Update State Representation Network

[0033]

[0034] In equation (2), It is a representation of the current state. It represents the next state; the first term aligns the state change with the skill direction, and the second term constrains the magnitude of the state representation change; alignment term. Constraints that ensure state changes align with skill direction To prevent the state representation from changing too much or too little, where ∈ is a preset threshold and λ is a Lagrange multiplier;

[0035] Secondly, the policy network π(a|s,z) and the Q-function network are updated using a flexible actor-critic algorithm; where the reward function is set as an intrinsic reward function based on changes in state representation. Essentially, it is the projection of state representation changes onto the skill direction, and the exploration effect is measured by the dot product of the skill vector and the state change; the cumulative reward of the trajectory adopts the standard discount form R = ∑γ t r t The discount factor γ∈(0.9, 0.99) controls the importance of long-term returns; a policy network π(a|s, z) and two Q-function networks Q are maintained. 1(s，a,z) Q 2(s,a,z) The two Q-function networks are updated by minimizing the Bellman error based on the current reward and the value of the next state; the policy network is updated by policy entropy regularization and maximizing the expected cumulative reward, where the entropy regularization term is controlled by a temperature coefficient.

[0036] Secondly, embodiments of the present invention also provide an unsupervised reinforcement learning system based on a world model extension, the system comprising: a historical data acquisition module, a world model construction module, a world model training module, a learning model construction module, an initialization module, an interactive data generation module, a buffer module, a learning model training module, and a reinforcement learning control module; wherein,

[0037] The historical data acquisition module is used to acquire the historical offline dataset D of the interaction between the trained expert strategy and the environment;

[0038] The world model construction module is used to construct a world model neural network;

[0039] The world model training module is used to train the world model neural network based on the historical offline dataset D, so as to obtain a mature world model neural network.

[0040] The learning model building module is used to build unsupervised reinforcement learning models;

[0041] The initialization module is used to set the initial current prediction state of the unsupervised reinforcement learning model;

[0042] The interactive data generation module is used to input the current predicted state and skill vector z into the policy network of the unsupervised reinforcement learning model to obtain the action; and input the current predicted state and action into the mature world model neural network to obtain the next predicted state; and determine whether the maximum step size or the termination state has been reached; if not, the next predicted state is used as the current predicted state, and the next action and the next predicted state are obtained; if so, the complete trajectory data is obtained and stored in the playback buffer.

[0043] The buffer module is used to set the playback buffer;

[0044] The learning model training module is used to randomly sample a portion of trajectory data from the playback buffer to train the unsupervised reinforcement learning model and obtain a mature unsupervised reinforcement learning model.

[0045] The reinforcement learning control module is used to input the current environmental state and the agent's skill vector z into a mature unsupervised reinforcement learning model, output the agent's action parameters, and control the agent to perform corresponding actions according to the action parameters.

[0046] The technical solutions provided in the embodiments of the present invention have the following beneficial effects:

[0047] The unsupervised reinforcement learning method and system based on world model extension provided in this invention constructs a more robust dynamic environment prediction system by integrating a world model architecture and differentiated prior training. The innovative skill learning framework ensures that learned skills have clear semantic meaning through explicit alignment of state representation changes with skill vectors. A dynamically adjusted constraint mechanism balances exploration efficiency and stability. This invention enhances the generalization ability and controllability of skills while improving learning efficiency. It is suitable for intelligent systems that need to learn interpretable and controllable skills from limited offline data, and has significant application value in fields such as robot operation and autonomous driving.

[0048] Of course, implementing any product or method of the present invention does not necessarily require achieving all of the advantages described above at the same time. Attached Figure Description

[0049] 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.

[0050] Figure 1 This is a flowchart of the unsupervised reinforcement learning method based on world model extension as described in the embodiments of the present invention;

[0051] Figure 2 This is a schematic diagram of the world model structure described in an embodiment of the present invention. Detailed Implementation

[0052] In their in-depth research of existing unsupervised reinforcement learning techniques, the inventors of this application discovered three key technical bottlenecks in skill learning and generalization. First, the skill learning process is disconnected from the dynamic characteristics of the environment, leading to significant instability in actual interactions, particularly a sharp decline in performance when facing unseen state spaces. Second, existing methods lack a systematic mechanism to ensure sufficient skill coverage in the state space, often resulting in certain critical regions remaining unexplored. Finally, insufficient modeling of the correlation between skill representation and state changes hinders precise skill control and transfer. These problems are particularly prominent in practical applications such as robot control and autonomous driving, severely restricting the industrial application of unsupervised reinforcement learning techniques.

[0053] Through extensive experimental analysis and theoretical derivation, the inventors recognized that the root cause of these technical bottlenecks lies in the failure of existing methods to effectively integrate dynamic environmental modeling with the skill learning process. Specifically, traditional methods either rely excessively on the static distribution of offline data, resulting in a lack of adaptability in the learned skills; or they lack reliable predictive ability for state transitions in online interactions, leading to low exploration efficiency. Furthermore, existing skill learning objectives often focus only on the statistical characteristics of the skill itself, neglecting the actual state changes that occur after the skill is executed. This disconnect makes it difficult for the learned skills to play their intended role in actual tasks.

[0054] Based on these important findings, this invention proposes an unsupervised reinforcement learning method based on world model extension. The implementation of this invention will be described in detail below with reference to the accompanying drawings and specific embodiments, including key aspects such as world model construction and training, reinforcement learning framework implementation, and core skill learning algorithms.

[0055] It should be noted that similar reference numerals and letters in the following figures indicate similar items; therefore, once an item is defined in one figure, it does not need to be further defined and explained in subsequent figures. In the description of this invention, the terms "first," "second," "third," "fourth," etc., are used only to distinguish descriptions and should not be construed as indicating or implying relative importance.

[0056] like Figure 1 As shown in the embodiment of the present invention, an unsupervised reinforcement learning method based on world model extension aims to solve the problem of insufficient generalization in offline dataset training, and includes the following steps:

[0057] Step S1: Obtain the historical offline dataset D of the trained expert policy and its interaction with the environment.

[0058] In this step, the collected historical offline dataset D is used to train the world model, and the data structure includes the environmental state and agent actions at different times.

[0059] Step S2: Construct the world model neural network.

[0060] In this step, an ensemble model architecture is used to construct a world model neural network; such as Figure 2 As shown, the world model neural network comprises n sub-models with identical structures. To enable each sub-model to fit the dataset distribution in different ways, each sub-model consists of a main network and a prior network, with the current environment state s serving as the input to both networks. t and the action a performed by the intelligent agent t The output is the predicted difference Δs between the predicted next state and the current state. t The parameters of the prior network for each sub-model are fixed and different, and its role is to act as an "anchor point" to guide the main network to update in different directions.

[0061] In this step, the sub-model employs a Multilayer Perceptron (MLP) structure. The integrated model architecture utilizes a parallel computing architecture, and the constructed world model neural network supports multiple sub-models simultaneously predicting states. The value of n is typically set between 5 and 20 depending on the task complexity. The MLP architecture contains three hidden layers with 256, 256, and 128 neurons respectively, and uses the Swish activation function to balance training stability and expressive power. The model's input is the current state. and actions Where d and m represent the dimensions of the state and action, respectively. The model output is the predicted state change, i.e., the prediction difference. The corresponding next state is s t+1 .

[0062] Step S3: Based on the historical offline dataset D, train the world model neural network to obtain a mature world model neural network.

[0063] In this step, training the world model neural network specifically includes: calculating the difference between the next state and the current state at each time step based on the environmental state and agent actions at different times in the historical offline dataset D, merging the difference into the dataset to construct a new dataset D*, and dividing D* into a training set, a validation set, and a test set; then training the sub-model main network in the world model neural network based on the training set in D*, and validating and testing it through the validation set and the test set respectively.

[0064] The world model neural network is trained using a differentiated prior method.

[0065] During training, the prior network of each sub-model does not participate in direct training; its parameters remain fixed and different, but it makes predictions based on the input values ​​to obtain the prediction difference of the prior network. The main network of each sub-model directly participates in training. After updating the main network parameters in each training epoch, it makes predictions based on the input values ​​to obtain the prediction difference of the main network. The output of the sub-model is a weighted sum of the prediction differences output by the main network and the prior network. Through this weighted summation, the prior network of each sub-model participates in the training process of the main network. During training, the main network needs to adjust its parameters to fit the weighted target, which includes fixed prior predictions. Therefore, the main network is forced to learn dynamic characteristics that complement or adjust the prior network, rather than fitting the data completely freely. In this way, different sub-models will ultimately learn different dynamic characteristics due to their different priors, thus achieving differentiation.

[0066] During training, the composite loss function used comprises four key components: a weighted mean squared error loss term to minimize state prediction bias, a log-variance loss term to quantify prediction uncertainty, a model parameter regularization term to prevent overfitting, and a log-variance range constraint term to ensure numerical stability. The mathematical expression for this loss function L is:

[0067]

[0068] In equation (1), N is the number of models in an ensemble, i represents the i-th ensemble model; B is the batch size, j represents the j-th batch; W is the output dimension, k represents the k-th output dimension; mean i,j,k logvar represents the mean of the predictions in the k-th dimension of the i-th ensemble model in the j-th batch; i,j,k Targets represents the mean, logarithm, and variance of the predictions in the k-th dimension for the i-th ensemble model in the j-th batch. j,k Let λ1 represent the target value in the j-th batch of the i-th ensemble model, where λ1 is the weight decay coefficient and θ is the neural network weight.

[0069] Preferably, the training process is terminated when the model performance does not significantly improve on the validation set for 5 consecutive epochs.

[0070] The detailed calculation process of the composite loss function is as follows: First, calculate the mean and logarithmic variance of the predicted output of each sub-model, where the logarithmic variance represents the logarithmic variance of the prediction uncertainty. Obtain the inverse variance through exponential operations, and then calculate the weighted mean squared error term:

[0071]

[0072] In this method, adaptive weighting is performed on prediction errors with different uncertainties, and a variance loss term is calculated at the same time. The model further includes a regularization term, which is the parameter L2 norm decay λ1||θ|| 2 . The final loss is the weighted sum of each term. Preferably, the Adam optimizer is used in the training process, the initial learning rate is set to 3e-4, and a cosine annealing scheduling strategy is adopted. The batch size is set between 256 and 1024 according to the available video memory, and the number of training epochs is usually 100 to 200.

[0073] Step S4: constructing an unsupervised reinforcement learning model, and setting an initial current prediction state.

[0074] In this step, the constructed unsupervised reinforcement learning model includes three core components: a state representation network a policy network π, and a Q-function network. Wherein, the state representation network adopts an MLP structure including two hidden layers, the input of which is the original state and the output thereof is a low-dimensional representation with k<<d, and the network uses the ReLU activation function to ensure training stability. The state representation network is configured to extract key features of the state. The policy network π also adopts an MLP structure including two hidden layers, the input of which is the state representation and a skill vector z, and the output thereof are the mean and variance parameters of the action distribution. To constrain the action range, a tanh activation function is used in the output layer, and the reparameterization trick is adopted to implement a differentiable stochastic policy. The Q-function network has a similar structure to the policy network, the input of which is the state representation the skill vector z, and an action a, and the output thereof is a state-skill value estimation. To improve learning stability, two target value networks are maintained in this embodiment, and the corresponding target network update rate τ is set to 0.005. All network parameters are initialized by the Xavier method, the Adam optimizer is selected, the initial learning rate is uniformly set to 3e-4, and the gradient clipping technology (with a threshold set to 1.0) is adopted to prevent gradient explosion.

[0075] Step S5: inputting the current prediction state and the skill vector z into the policy network of the unsupervised reinforcement learning model to obtain an action.

[0076] This step is a process of interaction between the world model neural network and the unsupervised reinforcement learning model to obtain training data of the unsupervised reinforcement learning model. Wherein, when the first interaction is performed, the current prediction state is the set initial current prediction state. Before the start of each round of interaction, a skill vector z is sampled from a preset standard Gaussian distribution N(0,1) as a signal for guiding the behavior of the agent in the current round. The skill vector z is input into the policy network of the unsupervised reinforcement learning model to obtain the corresponding action.

[0077] Step S6: Input the current predicted state and action into the mature world model neural network to obtain the next predicted state; and determine whether the maximum step size or the termination state has been reached; if not, take the next predicted state as the current predicted state and return to step S5; if yes, obtain the complete trajectory data, store it in the playback buffer, and proceed to step S7.

[0078] This step continues the interaction and training data collection process. From the n sub-models of the trained ensemble world model, one sub-model is randomly selected with equal probability as the simulation environment for the current round of interaction. In the current interaction round, the current predicted state and action are input into the selected sub-model, and the following interaction loop is performed at time steps t = 0, 1, 2, ..., T-1 (T is the preset maximum step size or until the termination state is reached):

[0079] Step S61, State Input and Action Generation: Input the current predicted state s t The skill vector z is input into the policy network π(a|s,z), and the policy network outputs the action a. t .

[0080] Step S62, World Model Prediction: The current predicted state s t and the generated action a t The input is fed into the selected world model sub-model. This sub-model outputs the predicted difference Δs between the predicted next state and the current predicted state. t .

[0081] Step S63, Calculate the next predicted state: Calculate the next predicted state predicted by the world model:

[0082] Step S64, Prediction Uncertainty Assessment: Calculate the uncertainty of the currently selected sub-model for the input (s) t ,a t The uncertainty of the prediction, u t The uncertainty u t The calculation method is as follows: (s) t ,a t Input the data into all n sub-models and calculate the predicted Δs. t The variance of the result, and use this variance as u t This variance reflects the integration discrepancies.

[0083] Step S65, Data Filtering and Storage: Determine Uncertainty u t Is it below the preset confidence threshold δ (δ>0)? If u t If the value is less than δ, the current prediction is considered reliable, and the state transition tuple (s) is set. t ,a t , z) Store in the playback buffer; if u t If the value is greater than or equal to δ, then discard the tuple to avoid low-quality data polluting the training.

[0084] Step S66, State Update: Update the next predicted state The current predicted state s for the next time step t+1 Determine if the maximum step size has been reached. If not, return to step S61. If yes, end the interaction and obtain the filtered, high-quality, complete trajectory data.

[0085] The replay buffer mentioned in this step also manages the replay buffer pool: when the replay buffer pool reaches its capacity limit, a buffer pool management operation is performed, that is, the oldest batch data is removed.

[0086] Step S7: Randomly sample a portion of trajectory data from the playback buffer to train the unsupervised reinforcement learning model and obtain a mature unsupervised reinforcement learning model.

[0087] In this step, the unsupervised reinforcement learning model is trained. The loss value of the unsupervised reinforcement learning model is calculated based on the data in the replay buffer pool, and the parameters are updated. It is determined whether the maximum number of training rounds has been reached. If not, return to step S5. If it has been reached, the training is completed, and a mature unsupervised reinforcement learning model is obtained.

[0088] The parameter update process includes a state representation network. The parameters of the three components, namely the policy network π and the Q-function network, are updated.

[0089] First, a loss function is learned using state representation. Update State Representation Network

[0090]

[0091] In equation (3), It is a representation of the current state. It represents the next state; the first term aligns the state change with the skill direction, and the second term constrains the magnitude of the state representation change; alignment term. Constraints that ensure state changes align with skill direction To prevent excessively large or small changes in state representation, ∈ is a preset threshold (recommended value is 0.1 to 0.3). λ is the Lagrange multiplier, updated using gradient descent.

[0092] Secondly, the flexible actor-critic (SAC) algorithm is used to update the policy network π(a|s,z) and the Q-function network. The reward function is set as an intrinsic reward function based on changes in state representation. Essentially, it's the projection of state representation changes onto the skill direction, and the exploration effect is measured by the dot product of the skill vector and the state change. The cumulative trajectory reward uses the standard discount form R = ∑γ t r t The discount factor γ∈(0.9, 0.99) controls the importance of long-term rewards. This reward design makes the agent inclined to produce state changes consistent with the current skill direction, thereby achieving skill-oriented exploration. A policy network π(a|s, z) and two Q-function networks Q are maintained. 1(s,a，z) Q 2(s，a，z) The two Q-function networks are updated by minimizing the Bellman error based on the current reward and the value of the next state; the policy network is updated by policy entropy regularization and maximizing the expected cumulative reward, where the entropy regularization term is controlled by a temperature coefficient. During training, the temperature coefficient is automatically adjusted to maintain a preset target policy entropy level to preserve the policy's exploratory nature.

[0093] The final trained policy network can generate corresponding goal-oriented actions based on the input skill vector z, achieving diverse skill control.

[0094] Step S8: Input the current environment state and the agent's skill vector z into the mature unsupervised reinforcement learning model, output the agent's action parameters, and control the agent to perform corresponding actions according to the action parameters.

[0095] In this step, the intelligent agent includes the Ant quadruped robot system simulated by Mujoco, an autonomous driving system, etc. Taking the robot as an example, the trained unsupervised reinforcement learning model generates corresponding goal-oriented actions based on the input skill vector z and the current environmental state parameters of the quadruped robot, thereby achieving diverse skill control.

[0096] Based on the same idea, this invention also provides an unsupervised reinforcement learning system based on a world model extension. The system includes: a historical data acquisition module, a world model construction module, a world model training module, a learning model construction module, an initialization module, an interactive data generation module, a buffer module, a learning model training module, and a reinforcement learning control module; wherein...

[0097] The historical data acquisition module is used to acquire the historical offline dataset D of the interaction between the trained expert strategy and the environment;

[0098] The world model construction module is used to construct a world model neural network;

[0099] The world model training module is used to train the world model neural network based on the historical offline dataset D, so as to obtain a mature world model neural network.

[0100] The learning model building module is used to build unsupervised reinforcement learning models;

[0101] The initialization module is used to set the initial current prediction state of the unsupervised reinforcement learning model;

[0102] The interactive data generation module is used to input the current predicted state and skill vector z into the policy network of the unsupervised reinforcement learning model to obtain the action; and input the current predicted state and action into the mature world model neural network to obtain the next predicted state; and determine whether the maximum step size or the termination state has been reached; if not, the next predicted state is used as the current predicted state, and the next action and the next predicted state are obtained; if so, the complete trajectory data is obtained and stored in the playback buffer.

[0103] The buffer module is used to set the playback buffer;

[0104] The learning model training module is used to randomly sample a portion of trajectory data from the playback buffer to train the unsupervised reinforcement learning model and obtain a mature unsupervised reinforcement learning model.

[0105] The reinforcement learning control module is used to input the current environmental state and the agent's skill vector z into a mature unsupervised reinforcement learning model, output the agent's action parameters, and control the agent to perform corresponding actions according to the action parameters.

[0106] In this embodiment, each module is implemented using a processor, with additional memory added as needed for storage. The processor can be, but is not limited to, a microprocessor (MPU), a central processing unit (CPU), a network processor (NP), a digital signal processor (DSP), an application-specific integrated circuit (ASIC), a field-programmable gate array (FPGA), other programmable logic devices, discrete gates, transistor logic devices, discrete hardware components, etc. The memory can include random access memory (RAM) or non-volatile memory (NVM), such as at least one disk storage device. Optionally, the memory can also be at least one storage device located remotely from the aforementioned processor.

[0107] In the above embodiments, implementation can be achieved, in whole or in part, through software, hardware, firmware, or any combination thereof. When implemented in software, it can be implemented, in whole or in part, as a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, all or part of the processes or functions described in the embodiments of the present invention are generated. The computer can be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions can be stored in a computer-readable storage medium or transmitted from one computer-readable storage medium to another. For example, the computer instructions can be transmitted from one website, computer, server, or data center to another website, computer, server, or data center via wired (e.g., coaxial cable, fiber optic, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) means.

[0108] It should also be noted that the unsupervised reinforcement learning system based on world model extension described in this embodiment corresponds to the unsupervised reinforcement learning method based on world model extension. The description and limitations of the method also apply to the system, and will not be repeated here.

[0109] As can be seen from the above technical solutions, the unsupervised reinforcement learning method and system based on world model extension provided by the embodiments of the present invention effectively solves key problems in existing technologies such as insufficient skill generalization and poor environmental adaptability. It provides a reliable technical foundation for the practical application of unsupervised reinforcement learning in fields such as robot control and intelligent manufacturing. It is particularly noteworthy that the world model integration architecture and skill learning algorithm of the present invention have good versatility and can be easily transferred to different types of decision-making and control tasks, demonstrating broad application prospects.

[0110] The above description is merely a preferred embodiment of the present invention and an explanation of the technical principles employed, and is not intended to limit the scope of the claimed invention, but merely to illustrate preferred embodiments of the invention. Those skilled in the art should understand that the scope of the invention is not limited to the specific combinations of the above-described technical features, but should also cover other technical solutions formed by arbitrary combinations of the above-described technical features or their equivalents without departing from the inventive concept. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without inventive effort are within the scope of protection of the present invention.

Claims

1. An unsupervised reinforcement learning method based on an extended world model, characterized in that, Includes the following steps: Step S1: Obtain the historical offline dataset D of the trained expert policy and its interaction with the environment; Step S2: Construct the world model neural network; Step S3: Based on the historical offline dataset D, train the world model neural network to obtain a mature world model neural network; Step S4: Construct an unsupervised reinforcement learning model and set the initial current prediction state; Step S5: Input the current predicted state and skill vector z into the policy network of the unsupervised reinforcement learning model to obtain the action; Step S6: Input the current predicted state and action into the mature world model neural network to obtain the next predicted state; and determine whether the maximum step size or the termination state has been reached; if not, take the next predicted state as the current predicted state and return to step S5; if yes, obtain the complete trajectory data, store it in the playback buffer, and proceed to step S7. Step S7: Randomly sample a portion of trajectory data from the playback buffer to train the unsupervised reinforcement learning model and obtain a mature unsupervised reinforcement learning model. Step S8: Input the current environment state and the agent's skill vector z into the mature unsupervised reinforcement learning model, output the agent's action parameters, and control the agent to perform corresponding actions according to the action parameters; The unsupervised reinforcement learning model constructed in step S4 includes at least three components: a state representation network φ, a policy network π, and a Q-function network; wherein, The state representation network φ adopts a multi-layer perceptron (MLP) structure, comprising two hidden layers, and its input is the original state and its output is a low-dimensional representation , where k << d. This network uses the ReLU activation function to ensure training stability; the state representation network is used to extract key features of the state; The policy network π adopts an MLP structure, containing two hidden layers. The inputs are the state representation φ(s) and the skill vector z, and the outputs are the mean and variance parameters of the action distribution. In order to constrain the action range, the output layer uses the tanh activation function and employs the reparameterization technique to realize a differentiable stochastic policy. The Q-function network adopts an MLP structure, containing two hidden layers. The input states are the state representation φ(s), the skill vector z, and the action a, and the output is the state-skill value estimate. Step S6 inputs the current predicted state and action into a mature world model neural network to obtain the next predicted state. Specifically, this includes: sampling a skill vector z from a pre-defined Gaussian distribution as the policy guidance signal for the current round; and then, based on the policy network... The policy network interacts with the world model neural network. Based on the current predicted state and the skill vector z, the policy network generates action a. The world model neural network receives action a and the current predicted state, and outputs the prediction difference. The difference between the current predicted state and the model's output prediction is added to obtain the next predicted state; finally, complete trajectory data is generated. And store it in the playback buffer pool.

2. The method according to claim 1, characterized in that, Step S2 involves constructing a world model neural network, including: An integrated model architecture is adopted to construct a world model neural network. The world model neural network contains n sub-models with identical structures, and each sub-model consists of a main network and a prior network. The input of both networks is the current environment state. Actions performed by the intelligent agent The outputs are the predicted difference between the next state and the current state. Furthermore, the parameters of the prior network for each sub-model are fixed but different; among them, the current state and actions d and m represent the dimensions of the state and action, respectively.

3. The method according to claim 2, characterized in that, The sub-model adopts a multilayer perceptron structure, and the multilayer perceptron structure contains three hidden layers; the integrated model architecture adopts a parallel computing architecture, and the constructed world model neural network supports multiple sub-models to perform state prediction simultaneously.

4. The method according to claim 2, characterized in that, Step S2 trains a world model neural network based on the historical offline dataset D, specifically including: Based on the environmental states and agent actions at different times in the historical offline dataset D, the difference between the next state and the current state at each time step is calculated. The difference is then incorporated into the dataset to construct a new dataset D*. D* is divided into a training set, a validation set, and a test set. The main network of the sub-model in the world model neural network is then trained based on the training set in D*, and validated and tested using the validation set and the test set, respectively.

5. The method according to claim 4, characterized in that, The training set based on D* is used to train the sub-model main network in the world model neural network, and a differential prior method is used to train the world model neural network. During training, the prior network of each sub-model does not participate in direct training. Its parameters are fixed and different, but it makes predictions based on the input values ​​to obtain the prediction difference of the prior network. The main network of each sub-model directly participates in training. After updating the parameters of the main network in each training round, it makes predictions based on the input values ​​to obtain the prediction difference of the main network. The output of the sub-model is a weighted sum of the prediction differences output by the main network and the prior network. During training, the composite loss function used consists of four components: a weighted mean squared error loss term to minimize state prediction bias, a log-variance loss term to quantify prediction uncertainty, a model parameter regularization term to prevent overfitting, and a log-variance range constraint term to ensure numerical stability.

6. The method according to claim 1, characterized in that, Step S7 trains the unsupervised reinforcement learning model, calculates the loss value of the unsupervised reinforcement learning model based on the data in the replay buffer pool, and updates the parameters; determines whether the maximum number of training rounds has been reached; if not, returns to step S5. If this is achieved, the training is complete, and a mature unsupervised reinforcement learning model is obtained.

7. The method according to claim 1, characterized in that, When updating parameters, a state representation network is included. Parameter updates for the three components: policy network (π), Q-function network, and policy network (Q-function network). The parameter update process includes: First, a loss function is learned using state representation. Update State Representation Network : （2） In equation (2), It represents the current state and the next state; the first term aligns the state change with the skill direction, and the second term constrains the magnitude of the state representation change; alignment term. Constraints that ensure state changes align with skill direction To prevent changes in state representation from being too large or too small, among which For the preset threshold, For Lagrange multipliers; Secondly, the flexible actor-critic algorithm is used to update the policy network and the Q-function network; where the reward function is set as an intrinsic reward function based on changes in state representation. Essentially, it is the projection of state representation changes onto the skill direction, and the exploration effect is measured by the dot product of the skill vector and the state change; the trajectory cumulative reward adopts a standard discount form, in which the discount factor controls the importance of long-term rewards; maintain the policy network. and two Q-function networks , The two Q-function networks are updated by minimizing the Bellman error based on the current reward and the value of the next state; the policy network is updated by policy entropy regularization and maximizing the expected cumulative reward, where the entropy regularization term is controlled by a temperature coefficient.

8. A system for an unsupervised reinforcement learning method based on world model extension according to any one of claims 1-7, characterized in that, The system includes: a historical data acquisition module, a world model construction module, a world model training module, a learning model construction module, an initialization module, an interactive data generation module, a buffer module, a learning model training module, and a reinforcement learning control module; wherein, The historical data acquisition module is used to acquire the historical offline dataset D of the interaction between the trained expert strategy and the environment; The world model construction module is used to construct a world model neural network; The world model training module is used to train the world model neural network based on the historical offline dataset D, so as to obtain a mature world model neural network. The learning model building module is used to build unsupervised reinforcement learning models; The initialization module is used to set the initial current prediction state of the unsupervised reinforcement learning model; The interactive data generation module is used to input the current predicted state and skill vector z into the policy network of the unsupervised reinforcement learning model to obtain the action; and input the current predicted state and action into the mature world model neural network to obtain the next predicted state; and determine whether the maximum step size or the termination state has been reached; if not, the next predicted state is used as the current predicted state, and the next action and the next predicted state are obtained; if so, the complete trajectory data is obtained and stored in the playback buffer. The buffer module is used to set the playback buffer; The learning model training module is used to randomly sample a portion of trajectory data from the playback buffer to train the unsupervised reinforcement learning model and obtain a mature unsupervised reinforcement learning model. The reinforcement learning control module is used to input the current environmental state and the agent's skill vector z into a mature unsupervised reinforcement learning model, output the agent's action parameters, and control the agent to perform corresponding actions according to the action parameters.

Images