Offline reinforcement learning sequence modeling method based on internal and external selection state spaces

Abstract

The invention belongs to the technical field of reinforcement learning, and particularly relates to an offline reinforcement learning sequence modeling method based on internal and external selection state spaces. In order to solve the problem of inaccuracy of model learning and the problems of complex Transforemer calculation and high cost in a sequence modeling process, the method comprises the following steps: constructing element dependency representation according to Markov attributes among trajectory sequence elements, modeling the element dependency representation by using an internal Mamba module, capturing a short-term causal relationship in a trajectory, and establishing an element dependency representation model according to Markov attributes among the trajectory sequence elements. And in combination with a state space model and off-line reinforcement learning, modeling is performed on an off-line reinforcement learning track sequence by utilizing Mama with linear complexity in an external Mama module, and a long-term time relationship among elements is captured.

Description

Technical Field

[0001] The present invention belongs to the technical field of reinforcement learning, and in particular relates to an offline reinforcement learning sequence modeling method based on internal and external selection state space. Background Art

[0002] Existing domestic operating systems, such as OpenEuler, are becoming more mature, but they still face problems such as a lack of applications, a small user base, and poor performance. The key to solving these problems is to establish an OpenEuler software adaptation and optimization platform, which uses reinforcement learning algorithms to optimize the performance of business execution running on the operating system, so that the software can run more smoothly in domestic operating systems.

[0003] As a solution to performance tuning, reinforcement learning is essentially a Markov decision process. Existing reinforcement learning algorithms can be divided into two types: online reinforcement learning and offline reinforcement learning. Online reinforcement learning collects data to train the agent by interacting with the environment in real time. However, since online reinforcement learning requires frequent interaction with the environment and requires a lot of exploration during the interaction, it is very expensive in practical applications. To this end, offline reinforcement learning provides a solution. Offline reinforcement learning uses pre-collected data sets to perform policy learning to train the agent. It does not require any interaction with the environment, so it is cheaper and safer, and more suitable for application in real-world scenarios.

[0004] The current mainstream offline reinforcement learning algorithms can be divided into three types: Q-Learning, imitation learning, and sequence modeling. The Q-Learning algorithm guides the decision-making process of the intelligent agent by learning the action value function that can be obtained by performing a certain action in a specific state. In the decision-making process, the action value function is continuously adjusted through the time difference (TD) so that the intelligent agent gradually finds the optimal strategy. For example, Conservative Q-Learning (CQL) proposed by Aviral Kumar et al. solves the problem of overestimation of value caused by the distribution shift between the behavior strategy and the learning strategy in the offline dataset by finding the lower bound of the true action value function and learning a conservative value function. However, estimating the value function in an offline setting may lead to unstable learning behavior and sensitivity to hyperparameters. Imitation learning learns effective strategies by simulating expert demonstrations from high-quality or near-optimal datasets in a supervised manner. For example, Behavior Cloning (BC) proposed by Michael Bain et al. maps the states in the dataset to their corresponding actions and trains by optimizing the difference in the action probability distribution between the expert behavior strategy and the learning strategy. However, imitation learning is highly dependent on expert datasets, while in practical applications, the number of high-quality expert datasets is limited, and the imitator cannot achieve higher performance than the expert dataset. In recent years, influenced by the Transformer's ability to capture dependencies between elements in long-distance sequences, researchers have tried to abstract offline reinforcement learning into a sequence modeling task, and use the Transformer's excellent long sequence modeling ability to solve decision-making problems in offline reinforcement learning. For example, the Decision Transformer model proposed by Lili Chen et al. sets the sum of future rewards (RTG) as the target return, models the historical information in the trajectory sequence, and predicts the next action to achieve the RTG goal by training an autoregressive Transformer model. Decision Transformer uses a target conditional strategy to convert offline reinforcement learning into a supervised learning task for actions, avoiding the use of approximate value functions and improving the performance and stability of offline reinforcement learning methods. However, Decision Transformer has a cost problem in offline reinforcement learning tasks. Its computational complexity grows quadratically with the increase in the length of the trajectory sequence, and requires significant computational and memory costs in the trajectory sequence modeling process.

[0005] In summary, the current offline reinforcement learning model can only process the historical information of trajectory elements in parallel during the sequence modeling process, without considering the causal relationship implied in the Markov properties between the underlying trajectory elements as the basis for decision-making, resulting in inaccuracy in model learning and the problem of complex and costly Transforemer calculations in the sequence modeling process. Summary of the invention

[0006] The present invention provides an offline reinforcement learning sequence modeling method based on internal and external selection state space in view of the above problems. The present invention solves the problems of high Transformer computational cost and inaccurate model prediction encountered in the current offline reinforcement learning sequence modeling method by nesting the internal Mamba module and the external Mamba module. The internal Mamba module models the element dependency representation constructed according to the Markov attribute in the time step dimension to capture the short-term causal relationship between the elements of the offline reinforcement learning trajectory; the external Mamba module is used to process the trajectory sequence composed of elements of four modes of action, state, RTG value and short-term causal relationship, and captures the long-term time relationship between the trajectory elements by modeling the trajectory sequence across time steps. Therefore, Mamba in Mamba, as an offline reinforcement learning trajectory sequence modeling model, can capture a wide range of dependencies and element information between trajectory elements, and has a powerful trajectory sequence modeling capability and better model generalization capability.

[0007] To achieve the above object, the present invention adopts the following technical solutions:

[0008] The present invention provides an offline reinforcement learning sequence modeling method based on internal and external selection state space, comprising the following steps:

[0009] Step 1: preprocess and sample the trajectory file to obtain a sub-trajectory sample data set;

[0010] Step 2, input the sub-trajectory sample dataset into the pre-trained Mamba in Mamba model through the data loader, and output the prediction of future actions. The Mamba in Mamba model combines the selection state space model and trajectory sequence modeling in the offline reinforcement learning task, including element dependency representation, internal Mamba module and external Mamba module. The element dependency representation is constructed according to the Markov property in the input trajectory sequence to better express the causal relationship between adjacent time step elements of the input trajectory sequence; the internal Mamba module captures the short-term causal relationship in the trajectory sequence from the element dependency representation; the external Mamba module models the interaction process of different time steps in the trajectory sequence under the condition of given short-term causal relationship, and captures the long-term temporal relationship between the elements of the trajectory sequence.

[0011] Furthermore, the nodes in the element dependency representation represent the three elements of the trajectory: state, action, and RTG, and the edges represent the Markov properties between different elements, specifically:

[0012] Given an RTG, state, and action sequence of T time steps, for the state s at time step tt ∈(B,K1), action a t ∈(B,K2), RTG R t ∈(B,K3) and the state s at time step t+1 t+1 ∈(B,K1), where B represents the batch size, K1, K2, and K3 represent s t 、a t , R t The size of the feature dimension;

[0013] According to the Markov property between the elements of the trajectory sequence, the element pairs with causal relationship are merged in the feature dimension to construct the element dependency representation with causal relationship. The element dependency representation contains three kinds of dependency relationships: (s t ,a t ) Dependencies between pairs of elements (s t ,R t ) Dependencies between pairs of elements (s t ,a t ,s t+1 ) Dependencies between pairs of elements

[0014] Furthermore, the basic structure of the internal Mamba module is Mambablock, which is divided into two stages: input processing and relationship modeling. In the input processing stage, the constructed element dependency representation is serialized as the input of the internal Mamba module. In the relationship modeling stage, Mambablock is used to model the element dependency representations at different time steps to capture the short-term causal relationship between elements in the trajectory sequence.

[0015] Furthermore, the specific process of the internal Mamba module is:

[0016] Represent element dependencies with Markov properties and Serialize it and get and

[0017] The serialized element dependency representation is projected to the input dimension D through the linear layer, and then the element dependency representation after linear projection is and Perform the addition operation to obtain the element-dependent embedding feature Emb at time step t t ∈(B,L,D), the process is shown in formula (1):

[0018]

[0019] Embed t Input into the internal Mamba module, use the Mamba block layer to model the relationship, and obtain the short-term causal relationship L between elements at time step t t ∈(B,L,D).

[0020] Furthermore, the internal Mamba module is composed of i Mambablock layers, wherein the operation performed in the i-th Mamba block is as shown in formula (2):

[0021]

[0022] Among them, RMSNorm represents the normalization operation of the feature, t represents the current time step, and i represents the layer number of the current Mamba block. represents the output feature of the Mamba block in the i-1th layer at time step t, Represents the output features of the i-th layer Mamba block at time step t.

[0023] Furthermore, the basic structure of the external Mamba module is Mambablock, which is divided into two stages: input processing and relationship modeling. In the input processing stage, the output of the internal Mamba module is added to the original trajectory sequence in time step order to achieve information supplementation of the trajectory sequence. In the relationship modeling stage, Mambablock is used to model the trajectory sequences of all time steps to capture the long-term temporal relationships in the trajectory sequences.

[0024] Furthermore, the specific process of the external Mamba module is:

[0025] The short-term causal relations L1, L2, ..., L at each time step output by the internal Mamba module T Arrange them in time step order as additional input to the original trajectory sequence;

[0026] Project the elements in the input trajectory sequence from the original dimension to the hidden layer dimension D through the linear layer, learn the features of elements of different modes, obtain the element embedding features of the corresponding mode, and then add the time step information to the element embedding features;

[0027] The embedded features of RTG, state, action and short-term causal relationship of T time steps are stacked sequentially in the time step dimension. Each time step contains elements of four modes. After normalization, the input trajectory sequence τ of the external Manba module is obtained. T ∈(B,4×L,D), as shown in formula (3):

[0028]

[0029] The normalized trajectory sequence τ T It is fed into the external Mamba module, where the Mamba block layer is used to model the trajectory sequence, capture long-term temporal relationships, and finally generate the predicted actions required for decision making.

[0030] Furthermore, the external Mamba module is composed of j Mamba Block layers, wherein the operation performed in the j-th Mamba Block layer is as shown in formula (4):

[0031]

[0032] Among them, RMSNorm represents the normalization operation of the feature, T represents the total number of time steps in the trajectory sequence, j represents the layer number of the current Mamba block, represents the output features of the j-1th layer Mambablock, Represents the output features of the j-th layer Mamba block.

[0033] Furthermore, the Mambablock includes two branches. The first branch performs linear projection, convolution, activation function processing and SSM layer processing on the input sequence; the second branch performs linear projection and activation function processing on the input sequence; and finally, the processing results of the two branches are merged, linearly projected and output. The specific process is as follows:

[0034] The shape of the input sequence X is (B, L, D), where B represents the batch size, L represents the sequence length, and D represents the embedding dimension size;

[0035] In the first branch of Mambablock, the input sequence X is projected into a hidden state X through a linear layer, X∈B,L,ED), and the hidden state feature X is processed using a 1D convolutional layer and Silu activation function, and then fed to the SSM layer;

[0036] In the SSM layer, the processed hidden state feature X is linearly projected again to obtain parameters B, C and step parameter Δ, B∈(B,L,N), C∈(B,L,N), Δ∈(B,L,ED), and then the parameters A and B are discretized using the step parameter Δ through zero-order hold to obtain the discretized parameters and The mapping of input sequence features to output sequence features is completed through the selective scanning of the Mamba block to obtain the processed features X, X∈(B,L,ED).

[0037] Furthermore, it also includes the training of the Mamba inMamba model, the specific process is as follows:

[0038] The model is trained by setting the number of training times and batch size. In each training iteration, the data of each batch in the data loader is input into the Mamba inMamba model to predict the state, action and RTG elements of the future time step in the input trajectory sequence. The model is trained in a self-supervised manner, and the average value of the mean square error between the predicted action results and the actions of the original input sequence is used as the loss value of the batch. After that, the optimizer optimization algorithm is used to update the model parameters.

[0039] Compared with the prior art, the present invention has the following advantages:

[0040] The Mamba in Mamba model of the present invention constructs element dependency representation based on the Markov properties between elements of the trajectory sequence, uses the internal Mamba module to model the element dependency representation, captures the short-term causal relationship in the trajectory, and improves the problem of inaccurate model learning. In addition, the Mamba in Mamba model combines the state space model and offline reinforcement learning, uses the linear complexity Mamba in the external Mamba module to model the offline reinforcement learning trajectory sequence, captures the long-term time relationship between elements, and solves the sequence modeling cost problem currently faced by offline reinforcement learning. The Mamba in Mamba model comprehensively models the potential short-term causal relationship and long-term time relationship in the trajectory sequence by nesting the internal Mamba module and the external Mamba module, can more effectively process complex trajectory data, and improve the accuracy and efficiency of the model. Experimental results show that the performance of the Mamba in Mamba model in the D4RL benchmark task is better than that of other offline reinforcement learning algorithms, demonstrating the application potential of the state space model in offline reinforcement learning. If the Mamba in Mamba model is applied to the performance tuning task of the domestic operating system, the software and application programs will perform more smoothly. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] Figure 1 This is the network architecture diagram of the Mamba in Mamba model;

[0042] Figure 2 The construction process for element dependency representation;

[0043] Figure 3 This is a schematic diagram of the structure of Mambablock;

[0044] Figure 4 The process of inputting the structure and parameterization of the SSM layer. DETAILED DESCRIPTION

[0045] In order to further illustrate the technical solution of the present invention, the present invention is further described below through embodiments.

[0046] The offline reinforcement learning sequence modeling method based on internal and external selection state space of this embodiment includes the following steps:

[0047] Step 1: preprocess and sample the trajectory files in the OpenAI Gym environment to obtain a sub-trajectory sample data set;

[0048] Step 2: Input the sub-trajectory sample dataset into the pre-trained Mamba in Mamba model through the data loader, and output the prediction of future actions. The Mamba in Mamba model combines the selection state space model and trajectory sequence modeling in the offline reinforcement learning task. Its architecture diagram is shown in the figure below: Figure 1 As shown, it includes element dependency representation, internal Mamba module and external Mamba module. The element dependency representation is constructed according to the Markov property in the input trajectory sequence to express the causal relationship between adjacent time step elements of the input trajectory sequence; the internal Mamba module captures the short-term causal relationship in the trajectory sequence from the element dependency representation; the external Mamba module models the interaction process of different time steps in the trajectory sequence given the short-term causal relationship to capture the long-term temporal relationship between the elements of the trajectory sequence.

[0049] Element dependency representation: Since there are Markov properties between the state, action, and future reward sum (RTG) elements of adjacent time steps in the trajectory sequence, that is, the action selection of the current time step depends on the state and RTG of the current time step, and the state generation of the current time step depends on the state and action of the previous time step. In order to model the Markov properties between the elements of the trajectory sequence, the present invention constructs an element dependency representation from the input trajectory and models the causal relationship between the elements of adjacent time steps, so as to help the model better handle the complex element relationship in the trajectory sequence during the sequence modeling process and improve the accuracy of the model prediction. The construction process of the element dependency representation is as follows: Figure 2 As shown, the nodes in the element dependency representation represent the three elements of state, action and RTG in the trajectory, and the edges represent the Markov properties between different elements.

[0050] Given a sequence of RTG, states, and actions of T time steps, for the state s at time step t t ∈(B,K1), action a t ∈(B,K2), RTG R t ∈(B,K3) and the state s at time step t+1 t+1∈(B,K1), where B represents the batch size, K1, K2, and K3 represent s t 、a t , R t The size of the feature dimension;

[0051] According to the Markov property between the elements of the trajectory sequence, the element pairs with causal relationship are merged in the feature dimension to construct the element dependency representation with causal relationship. The element dependency representation contains three kinds of dependency relationships: (s t ,a t ) Dependencies between pairs of elements (s t ,R t ) Dependencies between pairs of elements (s t ,a t ,s t+1 ) Dependencies between pairs of elements By constructing element dependency representations at different time steps, the short-term causal relationships in trajectory sequences can be modeled, thereby better utilizing the Markov properties in trajectory sequences in the sequence modeling process and improving the prediction accuracy and decision quality of the model.

[0052] Internal Mamba module: The basic structure is Mambablock, which is mainly used to capture the short-term causal relationship between elements in the trajectory sequence. It is divided into two stages: input processing and relationship modeling. In the input processing stage, the element dependency representation built based on the Markov property in the trajectory sequence is serialized as the input of the internal Mamba module. In the relationship modeling stage, Mamba block is used to model the element dependency representation at different time steps to capture the short-term causal relationship between elements in the trajectory sequence. Taking time step t as an example, the specific process is as follows:

[0053] Represent element dependencies with Markov properties and Serialize it and get and

[0054] The serialized element dependency representation is projected to the input dimension D through the linear layer, and then the element dependency representation after linear projection is and Perform the addition operation to obtain the element-dependent embedding feature Emb at time step t t ∈(B,L,D), the process is shown in formula (1):

[0055]

[0056] Embed t Input into the internal Mamba module, use the Mamba block layer to model the relationship, and obtain the short-term causal relationship L between elements at time step t t ∈(B,L,D).

[0057] The internal Mamba module consists of i Mambablock layers, where the operation performed in the i-th Mamba block is shown in formula (2):

[0058]

[0059] Among them, RMSNorm represents the normalization operation of the feature, t represents the current time step, and i represents the layer number of the current Mamba block. represents the output feature of the Mamba block in the i-1th layer at time step t, Represents the output feature of the Mamba block at time step t in the i-th layer. The dependency representation of the elements at time step t is modeled through the internal Mamba module to obtain the short-term causal relationship L between the elements at time step t t ∈(B,L,D). These short-term causal relationships in time steps will be used as supplementary information of the trajectory for the trajectory sequence modeling task in the external Mamba module.

[0060] External Mamba module: The structure of the external Mamba module is the same as that of the internal Mamba module. The difference is that the internal Mamba module models the element dependency representation, while the external Mamba module models the trajectory sequence. The basic structure of the external Mamba module is Mambablock, which is divided into two stages: input processing and relationship modeling. In the input processing stage, the output of the internal Mamba module is added to the original trajectory sequence in time step order to supplement the information of the trajectory sequence. In the relationship modeling stage, the Mamba block is used to model the trajectory sequence of all time steps to capture the long-term temporal relationship in the trajectory sequence. The specific process is:

[0061] First, the short-term causal relations L1, L2, ..., L at each time step output by the internal Mamba module are T Arrange them in time step order as additional input to the original trajectory sequence;

[0062] Then, the elements in the input trajectory sequence are projected from the original dimension to the hidden layer dimension D through a linear layer, the features of elements of different modalities are learned, the element embedding features of the corresponding modality are obtained, and the time step information is added to the element embedding features to ensure that the time information is not lost;

[0063] After that, the embedded features of RTG, state, action and short-term causal relationship of T time steps are sequentially stacked in the time step dimension. Each time step contains elements of four modes. After normalization, the input trajectory sequence τ of the external Manba module is obtained. T ∈(B,4×L,D), as shown in formula (3):

[0064]

[0065] Finally, the normalized trajectory sequence τ T The model is fed into the external Mamba module, and the Mambablock layer is used to model the trajectory sequence over a larger time span, capture long-term temporal relationships, and finally generate the predicted actions required for decision-making. Therefore, the Mamba inMamba model can more comprehensively understand and utilize the information in the trajectory by utilizing the nested two-layer Mamba structure, thereby improving the decision-making quality of offline reinforcement learning.

[0066] Mambablock: Mambablock is the infrastructure for internal and external Mamba modules. Its structure is as follows Figure 3 As shown. Mambablock consists of two branches. The first branch performs linear projection, convolution, activation function processing and SSM layer processing on the input sequence; the second branch performs linear projection and activation function processing on the input sequence; finally, the processing results of the two branches are merged, linearly projected and output. The specific process is as follows:

[0067] Assume that the shape of the input sequence X is (B, L, D), where B represents the batch size, L represents the sequence length, and D represents the embedding dimension size;

[0068] In the first branch of Mambablock, the input sequence X is projected into a hidden state X through a linear layer (Linear), where E is the expansion factor, the purpose of which is to enable the model to process higher-dimensional feature spaces and capture more detailed and complex features. The hidden state feature X is then processed using a 1D convolutional layer and a Silu activation function, and then fed to the SSM layer;

[0069] In the SSM layer (such as Figure 4As shown in the figure, the processed hidden state feature X is linearly projected again to obtain parameters B, C and step parameter Δ, B∈(B,L,N), C∈(B,L,N), Δ∈(B,L,ED). Therefore, the parameters in the SSM layer change dynamically with the input sequence, and the model can adaptively adjust its behavior according to the content of the input sequence. Then, the parameters A and B are discretized using the step parameter Δ through zero-order hold (ZOH) to obtain the discretized parameters and The mapping of input sequence features to output sequence features is completed through the selective scanning of the Mamba block to obtain the processed features X, X∈(B,L,ED).

[0070] In addition, the offline reinforcement learning sequence modeling method based on internal and external selection state space of this embodiment also includes the training and evaluation of the Mamba in Mamba model.

[0071] The present invention applies the Mamba model to the offline reinforcement learning sequence modeling task, inherits and develops the DecisionTransformer model by changing the network backbone structure. In terms of training and evaluation, Mamba in Mamba adopts a similar approach to Decision Transformer.

[0072] First, the trajectory data files in the OpenAI Gym environment are preprocessed and sampled to obtain the sub-trajectory sample dataset dataset. Then, the sub-trajectory sample dataset dataset is weighted randomly sampled according to the trajectory length by using the data loader DataLoader.

[0073] Then, the model model is trained by setting the number of training times max_iter and the batch size batch_size. In each training iteration, the data of each batch in the data loader DataLoader is input into the Mamba in Mamba model, and the state state_preds, action action_preds and RTG return_preds elements of the future time step in the input trajectory sequence are predicted. The model is trained in a self-supervised manner, and the average of the mean square error between the predicted action results action_preds and the actions actions of the original input sequence is used as the loss value train_loss of the batch, and then the optimizer optimization algorithm is used to update the model parameters.

[0074] Finally, the model is evaluated by setting the evaluation rounds num_eval_episodes and the target return target_return. In each evaluation round, the model's get_action method is used to obtain the next action action by interacting the initial state state with the environment, and the state state and reward reward of the next time step are obtained by executing the action action. The loop is executed until the maximum time step max_ep_len of the trajectory is reached, and the reward reward is accumulated to obtain the total reward episode_return obtained by the model in this round. Finally, the total returns obtained in multiple evaluation rounds are normalized and the average value is calculated to obtain the average normalized return of the model as the performance score score of the evaluation model. The model training and evaluation process is shown in Algorithm 1:

[0075]

[0076]

[0077] Example 2

[0078] The experimental environment of this embodiment is a server equipped with RTX3060Ti 8GB GPU.

[0079] In terms of training and evaluation data sets, this embodiment selects the benchmark data set D4RL, which is currently commonly used in offline reinforcement learning algorithms. D4RL provides a series of reinforcement learning benchmark data sets that are close to actual application scenarios. These data sets contain state-action-reward trajectory data collected from real robots or generated from simulation environments, and have the characteristics of proprioceptive observation and continuous action space. Since D4RL has advantages such as data richness, sample diversity, and environmental dynamic authenticity, it is more meaningful for training and evaluating the generalization ability and performance of offline reinforcement learning algorithms. In the D4RL benchmark, the present invention trains and evaluates the proposed Mamba in Mamba model in the simulation environment OpenAI Gym. Among them, Gym covers three different standard sports environments: HalfCheetah, Hopper and Walker, and each environment contains three different configurations: Medium-Expert, Medium and Medium-Replay, forming 9 environments of different sizes and qualities.

[0080] In terms of comparison methods, this paper selects the most representative methods among the three mainstream offline reinforcement learning algorithms, including Conservative q-learning (CQL) based on Q-Learning, behavioral cloning (BC) based on imitation learning, and Decision Transformer (DT) model based on sequence modeling.

[0081] In 9 experimental environments under the D4RL benchmark, the proposed Mamba in Mamba was trained and evaluated, and compared with the CQL, BC, and DT models. In the experiment, the average normalized return obtained by the model in the environment is used as the performance score, and the performance of the model in the environment is compared by comparing the performance scores obtained by different models in the same environment. The evaluation results are shown in Table 1. The data in the table represent the performance scores obtained by different models in the corresponding environment, among which the performance scores of CQL, BC, and DT are derived from the corresponding original papers. The score of 100 represents the score of the expert strategy, and the higher the score, the better the performance of the model in the corresponding environment.

[0082] The results in Table 1 show that the Mamba in Mamba model proposed in the present invention exhibits superior performance in the HalfCheetah-Medium-Expert, Hopper-Medium, Hopper-Medium-Replay, and Walker2d-Medium-Replay environments in D4RL, and near-optimal performance in the Hopper-Medium-Expert, Walker2dr-Medium-Expert, and Walker2d-Medium environments, especially when evaluated in the Hopper-Medium environment, the performance of the model is significantly improved. In addition, in all environments in the D4RL dataset, the average performance score of the Mamba inMamba model is the highest compared to other offline reinforcement learning models, and the Mamba in Mamba model achieves a 5.09% performance improvement compared to the suboptimal Decision Transformer. Therefore, Mamba in Mamba has better performance than the CQL, BC, and DT algorithms.

[0083] The performance score represents the average normalized return obtained by the model in the corresponding environment. The bold performance score represents the highest value in the current environment, and the underlined performance score represents the second highest value in the current environment. The results show that compared with CQL based on Q-learning, BC based on imitation learning, and DT based on sequence modeling, Mamba in Mamba shows better performance advantages.

[0084] Table 1 Model performance score results

[0085]

[0086] The above shows and describes the main features and advantages of the present invention. It is obvious to those skilled in the art that the present invention is not limited to the details of the above exemplary embodiments, and that the present invention can be implemented in other specific forms without departing from the spirit or essential features of the present invention. Therefore, no matter from which point of view, the embodiments should be regarded as exemplary and non-restrictive. The scope of the present invention is defined by the appended claims rather than the above description, and it is intended that all changes that fall within the meaning and scope of the equivalent elements of the claims are included in the present invention.

[0087] In addition, it should be understood that although the present specification is described according to implementation modes, not every implementation mode contains only one independent technical solution. This description of the specification is only for the sake of clarity. Those skilled in the art should regard the specification as a whole. The technical solutions in each embodiment may also be appropriately combined to form other implementation modes that can be understood by those skilled in the art.

Claims

1. An offline reinforcement learning sequence modeling method based on internal and external selection state space, characterized in that: The following steps are involved: Step 1: preprocess and sample the trajectory file to obtain a sub-trajectory sample data set; Step 2, input the sub-trajectory sample dataset into the pre-trained Mamba in Mamba model through the data loader, and output the prediction of future actions. The Mamba in Mamba model combines the selection state space model and trajectory sequence modeling in the offline reinforcement learning task, including element dependency representation, internal Mamba module and external Mamba module. The element dependency representation is constructed according to the Markov property in the input trajectory sequence to express the causal relationship between adjacent time step elements of the input trajectory sequence; the internal Mamba module captures the short-term causal relationship in the trajectory sequence from the element dependency representation; the external Mamba module models the interaction process of different time steps in the trajectory sequence under the given short-term causal relationship, and captures the long-term temporal relationship between the elements of the trajectory sequence.

2. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 1 is characterized in that: The nodes in the element dependency representation represent the three elements of the trajectory: state, action and RTG, and the edges represent the Markov properties between different elements, specifically: Given an RTG, state, and action sequence of T time steps, for the state s at time step t t ∈(B,K1), action a t ∈(B,K2), RTG R t ∈(B,K3) and the state s at time step t+1 t+1 ∈(B,K1), where B represents the batch size, K1, K2, and K3 represent s t 、a t , R t The size of the feature dimension; According to the Markov property between the elements of the trajectory sequence, the element pairs with causal relationship are merged in the feature dimension to construct the element dependency representation with causal relationship. The element dependency representation contains three kinds of dependency relationships: (s t ,a t ) Dependencies between pairs of elements (s t ,R t ) Dependencies between pairs of elements (s t ,a t ,s t+1 ) Dependencies between pairs of elements 3. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 2 is characterized in that: The basic structure of the internal Mamba module is Mambablock, which is divided into two stages: input processing and relationship modeling. In the input processing stage, the constructed element dependency representation is serialized as the input of the internal Mamba module. In the relationship modeling stage, Mambablock is used to model the element dependency representation at different time steps to capture the short-term causal relationship between elements in the trajectory sequence.

4. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 3 is characterized in that: The specific process of the internal Mamba module is: Represent element dependencies with Markov properties and Serialize it and get and The serialized element dependency representation is projected to the input dimension D through the linear layer, and then the element dependency representation after linear projection is and Perform the addition operation to obtain the element-dependent embedding feature Emb at time step t t ∈(B,L,D), the process is shown in formula (1): Embed t Input into the internal Mamba module, use the Mamba block layer to model the relationship, and obtain the short-term causal relationship L between elements at time step t t ∈(B,L,D).

5. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 3 or 4, characterized in that: The internal Mamba module is composed of i Mambablock layers, where the operation performed in the i-th Mambablock layer is as shown in formula (2): Among them, RMSNorm represents the normalization operation of the feature, t represents the current time step, and i represents the layer number of the current Mamba block. represents the output feature of the Mamba block in the i-1th layer at time step t, Represents the output features of the i-th layer Mambablock at time step t.

6. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 5 is characterized in that: The basic structure of the external Mamba module is Mambablock, which is divided into two stages: input processing and relationship modeling. In the input processing stage, the output of the internal Mamba module is added to the original trajectory sequence in time step order to supplement the information of the trajectory sequence. In the relationship modeling stage, the trajectory sequences of all time steps are modeled using Mamba block to capture the long-term temporal relationships in the trajectory sequence.

7. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 6 is characterized in that: The specific process of the external Mamba module is: The short-term causal relations L1, L2, ..., L at each time step output by the internal Mamba module T Arrange them in time step order as additional input to the original trajectory sequence; Project the elements in the input trajectory sequence from the original dimension to the hidden layer dimension D through the linear layer, learn the features of elements of different modes, obtain the element embedding features of the corresponding mode, and then add the time step information to the element embedding features; The embedded features of RTG, state, action and short-term causal relationship of T time steps are stacked sequentially in the time step dimension. Each time step contains elements of four modes. After normalization, the input trajectory sequence τ of the external Manba module is obtained. T ∈(B,4×L,D), as shown in formula (3): The normalized trajectory sequence τ T It is fed into the external Mamba module, where the Mamba block layer is used to model the trajectory sequence, capture long-term temporal relationships, and finally generate the predicted actions required for decision making.

8. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to claim 6 or 7, characterized in that: The external Mamba module is composed of j MambaBlock layers, where the operations performed in the j-th Mambablock are as shown in formula (4): Among them, RMSNorm represents the normalization operation of the feature, T represents the total number of time steps in the trajectory sequence, j represents the layer number of the current Mamba block, represents the output features of the j-1th layer Mambablock, Represents the output features of the j-th layer Mambablock.

9. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to any one of claims 3 to 8, characterized in that: The Mambablock includes two branches. The first branch performs linear projection, convolution, activation function processing and SSM layer processing on the input sequence; the second branch performs linear projection and activation function processing on the input sequence; finally, the processing results of the two branches are merged, linearly projected and output. The specific process is as follows: The shape of the input sequence X is (B, L, D), where B represents the batch size, L represents the sequence length, and D represents the embedding dimension size; In the first branch of Mambablock, the input sequence X is projected into a hidden state X through a linear layer, X∈B,L,ED), and the hidden state feature X is processed using a 1D convolutional layer and Silu activation function, and then fed to the SSM layer; In the SSM layer, the processed hidden state feature X is linearly projected again to obtain parameters B, C and step parameter Δ, B∈(B,L,N), C∈(B,L,N), Δ∈(B,L,ED), and then the parameters A and B are discretized using the step parameter Δ through zero-order hold to obtain the discretized parameters and The mapping of input sequence features to output sequence features is completed through the selective scanning of Mamba block, and the processed features X are obtained, X∈(B,L,ED).

10. The offline reinforcement learning sequence modeling method based on internal and external selection state space according to any one of claims 1 to 9, characterized in that: It also includes the training of the Mamba in Mamba model, the specific process is as follows: The model is trained by setting the number of training times and batch size. In each training iteration, the data of each batch in the data loader is input into the MambainMamba model to predict the state, action and RTG elements of the future time step in the input trajectory sequence. The model is trained in a self-supervised manner, and the average value of the mean square error between the predicted action results and the actions of the original input sequence is used as the loss value of the batch. Then the optimizer optimization algorithm is used to update the model parameters.

Images