Manipulation method of robot arm based on selective state space and model-based reinforcement learning

Abstract

The application provides a mechanical arm control method based on a selective state space and a model-based reinforcement learning, first, a world model is constructed according to a real environment, the world model processes a high-dimensional visual image under a third-person perspective, extracts image features and integrates sequence state-action information, a mechanical arm model makes decisions and learns according to the feature information extracted by the world model, can maximize sample efficiency under a lower interaction frequency, only uses a small amount of computing resources to enable the mechanical arm to learn an object grasping method, can improve the universality of the model-based reinforcement learning algorithm in different task scenarios, reduces training cost and deployment difficulty in an actual environment, and improves computing efficiency.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of reinforcement learning, the technical field of deep learning, and the technical field of manipulator control, and particularly relates to a manipulator control method based on selective state space and model-based reinforcement learning. BACKGROUND

[0002] Reinforcement learning is a paradigm of machine learning, which aims to let the agent learn the optimal policy through interaction with the environment. Reinforcement learning usually involves an agent, an environment, a state, an action, and a reward. The agent gets a reward after each interaction with the environment, tries different actions to get feedback (reward), and learns from these feedbacks to maximize the cumulative reward - return. With the rapid development of deep learning in recent years, deep reinforcement learning combines the powerful representation learning ability of deep neural networks with the decision optimization ability of reinforcement learning, enabling the agent to learn high-level abstract representations from raw environment data. In game scenarios such as StarCraft, Atari, and extremely difficult Go, deep reinforcement learning algorithms have repeatedly broken human records. In real-world application scenarios such as unmanned aerial vehicle racing and manipulator control, deep reinforcement learning has also achieved breakthrough performance improvement compared to traditional methods.

[0003] Model-free reinforcement learning algorithms such as PPO (Schulman J, Wolski F, Dhariwal P, et al. Proximal Policy Optimization Algorithms. [J]. CoRR, 2017, ) and SAC (Haarnoja T, Zhou A, Abbeel P, et al. Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor [C] / / Dy J, Krause A. Proceedings of Machine Learning Research: Proceedings of the 35th International Conference on Machine Learning) have achieved excellent results in many fields, but these results are often accompanied by millions of interactions. The extremely low sample efficiency limits the practical application of model-free reinforcement learning algorithms in environments with high safety performance requirements and high costs. To improve the sample efficiency of reinforcement learning algorithms, model-based reinforcement learning algorithms have emerged. Given that traditional model-free reinforcement learning methods often require a large number of interaction samples, PPO, SAC, and other model-free algorithms are usually trained in a simulation environment and then transferred to a real environment through a Sim To Real technique. Considering the complexity and randomness of the dynamics in the real environment, the performance of the strategy learned in the simulation environment in the actual scene is greatly limited. Moreover, traditional reinforcement learning algorithms such as PPO and SAC are difficult to handle high-dimensional input, which leads to a decrease in control accuracy in visual tasks.

[0004] Model-based reinforcement learning algorithms model the real environment, and the established model is called the world model. The agent interacts with the learned world model and trains the agent using the trajectories generated by the model. Compared with model-free reinforcement learning algorithms, model-based reinforcement learning algorithms alleviate the limitation of the number of interactions and significantly improve the sample efficiency. Dreamer algorithm (Hafner D, Lillicrap T P, Norouzi M, et al. Mastering Atari with discrete world models [C] / / International Conference on Learning Representations. 2021) is a typical model-based reinforcement learning algorithm in recent years. Dreamer uses recurrent neural networks as sequence models and establishes a complete variational inference process to build a recurrent state space model RSSM. Dreamer has shown strong learning ability in Atari games, DeepMind Control, Minecraft and other game environments. However, due to the natural sequential characteristics of RSSM, the world model is difficult to parallelize during training, resulting in slow training. In the task of manipulating a robotic arm, Dreamer directly learns the dynamics of the environment in the real environment and trains the agent's decision-making using the learned world model. Thanks to the randomness of the variational autoencoder, the strategy learned by the agent is more robust and can adaptively adjust the strategy according to the disturbance in the environment. Ultimately, Dreamer shows a training effect that is 5-6 times that of traditional model-based reinforcement learning methods such as PPO and Rainbow. However, due to the difficulty of parallelization during training, the entire training process often lasts nearly 10 hours, and the better training effect requires huge computing resources.

[0005] In recent years, some works use Transformer as the sequence model in the world model. This structure has shown strong ability in natural language processing, computer vision and other fields due to its unique attention mechanism, ability to establish long dependencies, and characteristics of large-scale parallel training. IRIS (Micheli V, Alonso E, Fleuret F. Transformers are sample-efficient world models [C] / / The Eleventh International Conference on Learning Representations. 2023) first encodes the image observation into 16 word vectors using VQ-VAE, and then uses a space-time Transformer structure to model the dynamic relationship between the image and the image. The world model built by IRIS has high precision, but because the space-time Transformer needs to perform self-attention operations on a large number of word vectors, the training speed is slow. STORM (Zhang W, Wang G, Sun J, et al. STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning [C] / / Thirty-seventh Conference on Neural Information Processing Systems. 2023) uses a single latent variable to represent the original image observation, and only uses two layers of random Transformer as the sequence model, which ensures the modeling accuracy while also providing high computational efficiency. It only takes 4.3 hours to complete the training of Atari 100K single game on a single NVIDIA GeForce RTX 3090. With the growth of the sequence, the inference time of the Transformer increases quadratically, which also causes the Transformer to be unable to provide the agent with long-term information, and the training of the agent is thus affected.

[0006] In addition to the Transformer, there are also scholars who use structured state space model (S4) to establish a world model. This structure has excellent long sequence modeling capability and can be trained in parallel through global convolution, and has achieved excellent performance in long sequence modeling tasks such as DNA. R2I (Samsami M R, Zholus A, Rajendran J, et al. Mastering Memory Tasks with World Models [C] / / The Twelfth International Conference on Learning Representations. 2024) and S4WM (Deng F, Park J, Ahn S. Facing Off World Model Backbones: RNNs, Transformers, and S4 [C] / / Thirty-seventh Conference on Neural Information Processing Systems. 2023)

[0007] The structured state space model shows excellent performance in environments that require long-term information to make decisions, such as MemoryMaze and Discount Chain. However, in simple environments that limit sample interactions, such as Atari100K, model-based reinforcement learning algorithms based on structured state space models are slightly inferior to Dreamer and other model-based reinforcement learning algorithms

[0008] The above-mentioned model-based reinforcement learning algorithms can improve the sample efficiency of reinforcement learning in multiple tasks, but there are still some problems, which are as follows:

[0009] 1) The performance of the agent in some environments is still very limited, and the average reward is low, which is far from the performance of humans;

[0010] 2) The computing efficiency is low, and the computing overhead and energy consumption are large in actual deployment. SUMMARY

[0011] To solve the above problems, the present application provides a mechanical arm control method based on selective state space and model-based reinforcement learning, which can maximize sample efficiency with fewer interactions, and only use a small amount of computing resources to make the mechanical arm learn the object grasping method, providing an efficient learning method for the mechanical arm.

[0012] A mechanical arm control method based on selective state space and model-based reinforcement learning, comprising the following steps:

[0013] S1: constructing a world model according to a real environment;

[0014] S2: interacting the robot arm model with the world model to obtain action sequences of the robot arm model under different environment observations;

[0015] S3: training the robot arm model by using the action sequences under different environment observations;

[0016] S4: controlling the robot arm to make corresponding actions in different real environments by using the trained robot arm model, so as to realize the control of the robot arm.

[0017] Further, the world model comprises an observation encoder q φ (z t |o t ), an image decoder an action information mixer m φ (z t ,a t ), a sequence model f φ (e 1:T ,x0), a dynamic predictor a reward predictor a termination predictor

[0018] The world model is constructed according to the real environment as follows:

[0019] S11: capturing the RGB image of the real environment in real time by a fixed image sensor to obtain the environment observation data o t of the current time t;

[0020] S12: sequentially reducing the dimension and randomly sampling the environment observation data o φ (z t |o t ) by the observation encoder q t , and representing the obtained random variable as a hidden space random variable z t by one-hot encoding;

[0021] S13: performing deconvolution on the hidden space random variable z t by the image decoder to obtain the reconstructed environment observation data

[0022] S14: concatenating the hidden space random variable z t and the action a t of the robot arm model under the environment observation data o t by the action information mixer m φ (z t ,a t) to map the concatenated data to state e t ;

[0023] S15: sequence model f based on Mamba architecture φ (e 1:T , x0) with state e t and the intermediate vector x t-1 at the previous time as input, output the recurrent state h t and the intermediate vector x t ;

[0024] S16: dynamic predictor receives recurrent state h t , outputs the latent random variable z

[0025] S17: reward predictor receives recurrent state h t and latent random variable z , outputs the reward r

[0026] S18: termination predictor receives recurrent state h t and latent random variable z , outputs the termination signal prediction value s

[0027] S19: construct a loss function according to the environment observation data reward termination signal prediction value , determine whether the loss function at the current time t is less than a set value, if yes, obtain the final world model, if no, enter step S110;

[0028] S110: adjust the parameters of the observation encoder q φ (z t |o t ), the image decoder action information mixer m φ (z t , a t ), the sequence model f φ (e 1:T , x0), the dynamic predictor reward predictor termination predictor , and then re-execute steps S11-S19 until the loss function is less than the set value.

[0029] Further, according to the environment observation data reward termination signal constructing a loss function as follows:

[0030]

[0031] where φ is an observation encoder q φ (z t |o t ), an image decoder an action information mixer m φ (z t ,a t ), a sequence model f φ (e 1:T ,x0), a dynamic predictor a reward predictor a termination predictor a set of parameters;

[0032] is an observation reconstruction loss for a current time t, is a reward prediction loss for the current time t, is a termination signal prediction loss for the current time t, is a dynamic prediction loss for the current time t, is a representation loss for the current time t, β1 is a weight of the dynamic prediction loss, and β2 is a weight of the representation loss.

[0033] Further, a calculation method of the observation reconstruction loss is as follows:

[0034]

[0035] where ||·||2 represents a calculation of the Euclidean distance;

[0036] a calculation method of the reward prediction loss is as follows:

[0037]

[0038] where twohot represents double-hot encoding, symlog(r t ) = sign(r t ) ln(|r t | + 1), and Reward Logits is a reward regression vector output by the reward predictor .

[0039] a calculation method of the termination signal prediction loss is as follows:

[0040]

[0041] Among them, c t This represents the actual value of the termination signal at the current time t;

[0042] Dynamically predicting losses The calculation method is as follows:

[0043]

[0044] Characterization loss The calculation method is as follows:

[0045]

[0046] Where sg() is the terminating gradient operator, z t+1 Let o be the latent space random variable at the next time step. t+1 For the environmental observation data captured by the image sensor at the next moment, KL[] represents the KL divergence. Indicates that the dynamic predictor is h t For the input predicted The distribution of .

[0047] Furthermore, by interacting the robotic arm model with the world model, the specific action sequences of the robotic arm model under different environmental observations are obtained as follows:

[0048] S21: Input the historical environmental observation data o0~o7 for 8 consecutive time steps and the historical actions a0~a7 performed by the robotic arm under the environmental observation data o0~o7 into the world model to obtain the initial intermediate vector x0, the loop state h0 and the latent space random variable z1;

[0049] S22: The world model passes h0 and z1 to the robotic arm model, and the robotic arm model outputs the decision action.

[0050] S23: The motion information mixer will combine z1 with... The concatenation is mapped to e1. ​​The sequence model takes e1 and x0 as input and outputs the current time step's loop state h1 and intermediate vector x1.

[0051] S24: Input h1 into the dynamic predictor and map h1 to a 1024-dimensional regression vector Logits;

[0052] S25: Divide the regression vector Logits into 32 groups of 32-dimensional vectors and perform softmax activation function operation to output 32 groups of prior distributions;

[0053] S26: Randomly sample 32 groups of prior distribution, and represent the sampling results by one-hot encoding, then concatenate the 32 groups of sampling results represented by one-hot encoding back to 1024 dimensions in sequence to obtain the hidden space random variable of the next time step predicted by the dynamic predictor

[0054] S27: Concatenate h1, to obtain the predicted reward and termination signal prediction value

[0055] S28: Repeat steps S22-S27 until T groups of trajectory data are obtained, wherein the T groups of trajectory data include the loop state h 0:T-1 , the predicted hidden space random variable the predicted reward and the predicted termination signal prediction value

[0056] Further, training the robot arm model using action sequences under different environment observations includes the following steps:

[0057] S31: Concatenate h t-1 and h t in the T groups of trajectory data respectively to obtain T states s t ;

[0058] S32: Obtain the value function V ψ (s t ) corresponding to each state s t according to the reward in the T groups of trajectory data respectively;

[0059] S33: Input each state s ψ into the robot arm model to output the decision action of the robot arm;

[0060] S34: Construct a training loss function of the robot arm model according to the value function V t (s t ) and the decision action, and determine whether the training loss function is less than a set value, if yes, the final robot arm model is obtained, if no, step S35 is entered;

[0061] S35: Change the parameters of the robot arm model, and re-execute steps S31-S34 until the training loss function is less than the set value.

[0062] Further, the training loss function of the robot arm model is as follows:

[0063]

[0064] Where θ is the parameter to be optimized in the policy function. Let ψ be the loss function corresponding to the policy function, and ψ be the parameter to be optimized in the value function. Here, is the loss function corresponding to the value function, sg() is the terminating gradient operator, B is the training batch, L is the length of the trajectory data used in each batch, and S is the set scaling factor. For state s t Select action The policy function is defined by HL() for Histogram encoding, max() for maximizing the value, and η for setting the weights. Entropy of the policy function, activation function Where V represents the value function V ψ () with s t The original output from the input is a 255-dimensional vector, v. j This represents the value of the j-th dimension of V. To estimate the true value function using time-series differencing, we have:

[0065]

[0066] Where λ is the set weight and γ is the attenuation coefficient.

[0067] Furthermore, the movement of the robotic arm is a seven-dimensional vector. The first three dimensions represent the coordinate changes of the robotic arm end effector along the x-axis, y-axis, and z-axis in the Cartesian coordinate system, respectively. The middle three dimensions represent the rotation angles of the robotic arm end effector with the x-axis, y-axis, and z-axis as rotation axes, respectively. The last dimension represents the opening and closing of the robotic arm end effector gripper.

[0068] Beneficial effects:

[0069] This invention provides a robotic arm control method based on selective state space and model-based reinforcement learning. First, a world model is constructed based on the real environment. The world model processes high-dimensional visual images from a third-person perspective, extracts image features, and integrates sequential state-action information. The robotic arm model then makes decisions and learns based on the feature information extracted from the world model. This maximizes sample efficiency with fewer interactions and allows the robotic arm to learn object grasping methods using minimal computational resources. It improves the generality of model-based reinforcement learning algorithms in different task scenarios, reduces training costs and deployment difficulty in real-world environments, and enhances computational efficiency. In other words, this invention utilizes the powerful sequential modeling capabilities of selective state space models and efficient parallel training methods, enabling the robotic arm to efficiently complete object grasping tasks under visual input after 5-6 hours of training. This provides a promising solution for intelligent production tasks, particularly suitable for the field of intelligent manufacturing. BRIEF DESCRIPTION OF DRAWINGS

[0070] Figure 1 is a schematic diagram of the world model of the present application when training in parallel;

[0071] Figure 2 is a schematic diagram of the world model of the present application when interacting with a robot arm model;

[0072] Figure 3 is a schematic diagram of the algorithm flow provided by the present application;

[0073] Figure 4 is a network information flow diagram provided by the present application. DETAILED DESCRIPTION

[0074] In order to enable persons skilled in the art to better understand the scheme of the present application, the technical solutions in the embodiments of the present application will be clearly and completely described below in conjunction with the drawings in the embodiments of the present application.

[0075] As shown in Figure 1 , a robot arm control method based on selective state space and model-based reinforcement learning comprises the following steps:

[0076] S1: constructing a world model according to a real environment; the world model comprises an observation encoder q φ (z t |o t ), an image decoder action information mixer m φ (z t ,a t ), sequence model f φ (e 1:T ,x0), dynamic predictor reward predictor termination predictor

[0077] Observation encoder: z t ~ q φ (z t |o t )

[0078] Image decoder:

[0079] Action information mixer: e t = m φ (z t ,a t )

[0080] Sequence model: h 1:T , x 1:T = f φ (e 1:T,x0)

[0081] Dynamic predictor:

[0082] Reward Predictor:

[0083] Termination Predictor

[0084] Among them, the observation encoder q φ (z t |o t ) and image decoder Composed of discrete distribution variational autoencoders;

[0085] like Figure 4 As shown, the world model constructed based on the real environment is as follows:

[0086] S11: Real-time RGB image capture of the real environment is obtained by a fixed image sensor, yielding environmental observation data at the current time t. t ;

[0087] It should be noted that the environmental observation data... t It is a 64×64-dimensional high-dimensional RGB image, captured in real time by a fixed image sensor; the high-dimensional environmental observation data has a dimension of over 3000, while the low-dimensional latent space random variables have a dimension of less than 1500.

[0088] S12: Observation encoder q composed of a convolutional neural network φ (z t |o t Environmental observation data t The resulting random variables are represented as latent space random variables z using one-hot encoding after dimensionality reduction and random sampling. t ;

[0089] Specifically, observe encoder q φ (z t |o t High-dimensional environmental observation data t The encoding is applied to a 1024-dimensional regression vector, Logits. Logits can be divided into 32 groups of 32-dimensional vectors, which are then processed by the softmax activation function to become 32 discrete arrays. Random sampling is performed using these 32 arrays, and the resulting random variables are represented using one-hot encoding. These 32 random variables are then concatenated sequentially to obtain the 1024-dimensional latent space random variable z. t ;z t It can be regarded as a high-dimensional observation. t The feature vectors extracted by convolutional neural networks and multilayer perceptrons contain o tAll important feature information; it should be noted that the discrete sampling step cannot pass the gradient, and the present application uses a "direct gradient passing" method for sampling (as shown in the following formula), which allows the gradient to pass from the decoder to the encoder through the probability probs. Where sg() represents the termination gradient operator:

[0090] z t = z t +probs-sg(probs)

[0091] S13: Image decoder composed of deconvolutional neural network Deconvolve the latent space random variable z t to obtain the reconstructed environment observation data

[0092] latent space random variable z t as input to the image decoder , Deconvolve z t to obtain the reconstructed environment data; as can be seen, the present application encodes high-dimensional environment observation data o t into low-dimensional latent space random variable z t , referred to as latent variable z t , z t contains feature information of environment observation.

[0093] S14: Concatenate the latent space random variable z t and the action a t of the robotic arm model under the environment observation data o t The action information mixer m φ (z t , a t ) composed of multiple layers of perception maps the concatenated data to state e t ;

[0094] It should be noted that the action of the robotic arm is a seven-dimensional vector, the first three dimensions represent the coordinate changes of the robotic arm end controller along the x, y, and z axes in the Cartesian coordinate system; the middle three dimensions represent the rotation angles of the robotic arm end controller with the x, y, and z axes as the rotation axes; the last dimension represents the opening and closing of the robotic arm end gripper.

[0095] After concatenating the latent variable z t and the action a t selected by the robotic arm in the current environment state and the current time step, the action information mixer m φ (z t , a t ) composed of multiple layers of perception maps to state et e t It can be regarded as a state variable that contains observation information and action information at the current time step.

[0096] S15: Sequence Model Based on Mamba Architecture φ (e 1:T ,x0) in state e t The intermediate vector x at the previous time step t-1 As input, output the loop state h at the current time t. t and intermediate vector x t ;

[0097] The sequence model f in this invention φ (e 1:T ,x0) with e t The sequence is taken as input, and the corresponding loop state h containing historical observation and action information is output. t Sequence model f φ (e 1:T The model x0 uses a selective state-space model, also known as Mamba. During world model training, because the complete sequence information can be received in advance, Mamba eliminates the need for single-step inference; intermediate states x0 are directly obtained. t The operations are performed internally within Mamba, requiring no input or output.

[0098] S16: By dynamic predictor Receive loop state h t Output the hidden space random variable at the next time step.

[0099] Specifically, this invention uses a dynamic predictor composed of a multilayer sensing mechanism. h t Mapped to a 1024-dimensional regression variable Logits; this invention divides the Logits into 32 groups of 32-dimensional vectors, and these 32 groups of vectors are processed by the softmax activation function to output 32 prior distributions. Random samples are taken from these 32 prior distributions, and the sampling results are represented using one-hot encoding. These 32 sampling results are then concatenated sequentially back to the 1024-dimensional array to obtain the latent variables for the next time step predicted by the dynamic predictor.

[0100] S17: By reward predictor Receive loop state h t and latent space random variables Output the reward at the current time t.

[0101] It should be noted that the reward For simple 0-1 signal, if the current time captures the article r t 1, otherwise 0; the present application uses a multi-layer perception mechanism to form a reward predictor The cycle state h t And the hidden variable of the next time step As input, the output dimension of the reward regression vector Reward Logits is 255. After the Reward Logits is operated by the activation function softmax, a set of weights The present application provides a set of reference vectors with equal intervals and a dimension of 255, [-20, -20+bin, …, 20-bin, 20], where the step size bin=40 / 254. The reference vector and the weight One-to-one correspondence, the vector and the weight are multiplied by matrix to obtain the predicted compressed reward, and the range of the compressed reward is [-20, 20]. The present application decompresses the compressed reward by symexp=sign(x)(e |x| -1) to obtain the predicted value of the current time step environment reward

[0102] S18: The termination predictor receives the cycle state h t And the hidden space random variable Outputs the termination signal prediction value of the current time t

[0103] It should be noted that the termination signal prediction value It is also a 0-1 signal, and the maximum time for a single interaction is limited to 6h. Before reaching the maximum time for a single interaction 1, after reaching the maximum interaction time 0, the training is over; the present application uses a multi-layer perception mechanism to form a termination predictor The cycle state h t And the hidden variable of the next time step predicted by the dynamic predictor As input, the termination signal of the current time step is predicted, and the output of the termination predictor

[0104] S19: According to the environment observation data Reward Termination signal prediction value The loss function is constructed to determine whether the loss function of the current time t is less than the set value. If yes, the final world model is obtained, and if no, step S110 is entered.

[0105] It should be noted that the world model in the application is trained in an end-to-end manner, and is trained in a gradient descent manner using an Adam optimizer, wherein the environment observation data reward termination signal constructing a loss function As follows:

[0106]

[0107] wherein φ is a parameter set of an observation encoder q φ (z t |o t ), an image decoder action information mixer m φ (z t ,a t ), a sequence model f φ (e 1:T ,x0), a dynamic predictor reward predictor termination predictor .

[0108] is an observation reconstruction loss of the current time t, wherein the superscript rec is the abbreviation of reconstruction; is a reward prediction loss of the current time t, wherein the superscript rew is the abbreviation of reward; is a termination signal prediction loss of the current time t, wherein the superscript con is the abbreviation of continuous; is a dynamic prediction loss of the current time t, wherein the superscript dyn is the abbreviation of dynamic; is a representation loss of the current time t, wherein the superscript rep is the abbreviation of representation; β1 is the weight of the dynamic prediction loss, and β2 is the weight of the representation loss; optionally, the weight β1=0.5, and the weight β2=0.1.

[0109] In each round of training, after the total loss function L(φ) is calculated, the Adam optimizer is used to update the world model parameters in a gradient descent manner.

[0110] Further, the observation reconstruction loss is used to update the encoder observation encoder and the image decoder, so as to improve the reconstruction accuracy of the world model, and the calculation method is as follows:

[0111]

[0112] wherein ||·||2 represents the calculation of the Euclidean distance;

[0113] Reward Prediction Loss Used to update the reward predictor and improve the accuracy of reward prediction, its calculation method is as follows:

[0114]

[0115] Where twohot represents two-hot encoding, symlog(r t ) = sign(r t )ln(|r t |+1), Reward Logits is a reward predictor. The output reward regression vector;

[0116] Termination signal prediction loss Used to update the termination predictor and improve the accuracy of termination signal prediction, its calculation method is as follows:

[0117]

[0118] Among them, c t This represents the actual value of the termination signal at the current time t;

[0119] Dynamically predicting losses The calculation method is as follows:

[0120]

[0121] Characterization loss The calculation method is as follows:

[0122]

[0123] Where sg() is the terminating gradient operator, z t+1 Let o be the latent space random variable at the next time step. t+1 For the environmental observation data captured by the image sensor at the next moment, KL[] represents the KL divergence. Indicates that the dynamic predictor is h t For the input predicted The distribution of .

[0124] It should be noted that the dynamic loss function With representation loss function Both are almost identical in form, both using the method of maximizing KL divergence for updating, and are only distinguished by the terminating gradient operator sg(); the dynamic loss function is used to improve the prediction accuracy of the dynamic predictor, while the characterization loss function is used to reduce the difficulty of dynamic prediction.

[0125] S110: Adjust the observation encoder q φ (z t|o t Image decoder Motion Information Mixer m φ (z t ,a t ), sequence model f φ (e 1:T (x0) Dynamic predictor Reward Predictor Termination Predictor The parameters are then set, and steps S11 to S19 are re-executed until the loss function is less than the set value.

[0126] After the world model completes training, it interacts with the robotic arm model via autoregression to generate imagined trajectory data, as shown in the attached figure. Figure 2 In the conceptualization phase, this invention uses Mamba's cyclic inference mode. At each inference step, the sequence model constructed using Mamba requires the input of the intermediate vector x from the previous time step. t Output the intermediate vector x for the next time step. t+1 For details, please refer to step S2.

[0127] S2: Interact the robotic arm model with the world model to obtain the action sequence of the robotic arm model under different environmental observations, such as... Figure 3 As shown, the specific steps include:

[0128] S21: Input the historical environmental observation data o0~o7 for 8 consecutive time steps and the historical actions a0~a7 performed by the robotic arm under the environmental observation data o0~o7 into the world model to obtain the initial intermediate vector x0, the loop state h0 and the latent space random variable z1;

[0129] S22: The world model passes h0 and z1 to the robotic arm model, and the robotic arm model outputs the decision action.

[0130] S23: The motion information mixer will combine z1 with... The concatenation is mapped to e1. ​​The sequence model takes e1 and x0 as input and outputs the current time step's loop state h1 and intermediate vector x1.

[0131] S24: Input h1 into the dynamic predictor and map h1 to a 1024-dimensional regression vector Logits;

[0132] S25: Divide the regression vector Logits into 32 groups of 32-dimensional vectors and perform softmax activation function operation to output 32 groups of prior distributions;

[0133] S26: randomly sample 32 groups of prior distribution, and use one-hot encoding to represent the sampling results, then concatenate the 32 groups of sampling results represented by one-hot encoding back to 1024 dimensions in sequence to obtain the hidden space random variable of the next time step predicted by the dynamic predictor

[0134] S27: h1, is passed to the reward predictor and the termination predictor to obtain the predicted reward and the termination signal prediction value

[0135] S28: repeat steps S22-S27 until T groups of trajectory data are obtained, wherein the T groups of trajectory data include the recurrent state h 0:T-1 , the predicted hidden space random variable the predicted reward the predicted termination signal prediction value As Figure 3 shown.

[0136] Wherein, 0:T-1 in the subscript represents from the 0th time step to the T-1th time step, for example, h 0:T-1 here represents the recurrent state from the 0th time step to the T-1th time step.

[0137] S3: train the robot arm model using the action sequence under different environment observations;

[0138] That is, step S3 trains the operation of the robot arm based on the trajectory data generated in step S2, and the goal is for the robot arm model to grasp more than a set value of objects within a limited time, and the robot arm is allowed to grasp as many objects as possible within a limited time by maximizing the cumulative reward;

[0139] The data used for training the robot arm model in the application is all derived from the interaction with the world model in step S2. The training of the robot arm uses an actor-critic framework, the actor is a parameterized policy function π θ which directly controls the robot arm action according to the feature information extracted after the environment information is arranged by the world model, and the critic is a parameterized value function V ψ which estimates the expected discounted cumulative reward under the current policy function, and the value function can guide the training of the policy function; the specific framework is as follows:

[0140] State:

[0141] Critic:

[0142] Actor: a t~π θ (a t |s t )

[0143] Among them, s t Will with h t-1 The process of splicing these elements together is called a state. π θ Here, θ is a parameterized policy function used to directly control the actions taken by the robotic arm at any given time step; V ψ This is a parameterized value function used to estimate the cumulative discount reward under the current strategy. Here, the subscript ψ represents the parameter of the value function, and γ = 0.975 is the decay coefficient. This represents the environmental reward generated by the world model at time step t+k. for To dynamically calculate the expectation of the robotic arm's movements in relation to the world model, this invention uses Monte Carlo sampling to approximate this expectation.

[0144] Specifically, the training method for the robotic arm model includes the following steps:

[0145] S31: Extract the trajectory data from group T respectively. with h t-1 By concatenating the states, we obtain T states s. t ;

[0146] S32: Rewards based on the trajectory data of group T Obtain each state s separately t The corresponding value function V ψ (s t );

[0147] S33: Separate each state s t Input a robotic arm model, output the robotic arm's decision action;

[0148] S34: According to the value function V ψ (s t ), construct the training loss function of the robotic arm model based on the decision action, and determine whether the training loss function is less than the set value. If yes, the final robotic arm model is obtained; if no, proceed to step S35.

[0149] The training loss function for the robotic arm model is as follows:

[0150]

[0151] Where θ is the parameter to be optimized in the policy function. Let ψ be the loss function corresponding to the policy function, and ψ be the parameter to be optimized in the value function. is the loss function corresponding to the value function, sg() is the terminal gradient operator, B is the training batch, L is the length of the trajectory data used at each batch, S is the set reduction multiple, is the state s t Selects the action of the policy function, HL() represents Histogram encoding, max() represents the maximum value, η represents the set weight, represents the entropy of the policy function, the activation function where V represents the value function V ψ () when s t is the original output, a vector of dimension 255, v j represents the value of the jth dimension of V, is estimated by means of temporal difference, and has:

[0152]

[0153] where λ is the set weight, and γ is the decay coefficient.

[0154] It should be noted that S limits the gradient size when updating the policy, representing the difference between the top 5% and the bottom 5% of the current all ; HL() represents Histogram encoding, where V ψ (s t ) uses a similar structure to the reward predictor, and the output is a Logits vector of dimension 255, but the encoding of the calibration data is different, the reward predictor uses TwoHot encoding, Histogram encoding; the Histogram encoding method can refer to Stop Regressing; updates the parameters of the value function V ψ by cross-entropy error, taking as the reference quantity; updates the policy function π θ by policy gradient, updating in the direction of maximizing cumulative reward, where is the entropy of the policy function, and the entropy of the policy function is expected to be greater than a certain value at the beginning of training to ensure sufficient exploration; during training, the robot learns in the direction of maximizing cumulative reward, and finally learns how to perform continuous translation, rotation, and gripper opening and closing actions based on the information z t and h t-1 extracted by the world model at the current time step, to complete the task of grasping the object; the robot model also uses the Adam optimizer during training, and is optimized by gradient descent.

[0155] S35: Change the parameters of the robot arm model, re-execute steps S31-S34 until the training loss function is less than a set value.

[0156] S4: Control the robot arm to make corresponding actions in different real environments using the trained robot arm model, to realize the control of the robot arm.

[0157] In summary, the present application provides a robot arm control method based on selective state space and model-based reinforcement learning. First, a world model is constructed between the task scene and the robot arm model. Then, the robot arm model and the world model interact to predict the trajectory data of the real environment in a self-recurrent manner. Finally, the robot arm model is trained based on the predicted trajectory data generated in step S2. Based on this, the present application has the following advantages:

[0158] When training the world model, the present application uses the Mamba hardware-aware acceleration algorithm to accelerate the training in a parallel scanning manner. Thanks to this, the training speed of the present application can be comparable to that of STORM, a model-based reinforcement learning algorithm based on Transformer. It only takes 4.5 hours to complete the training under Atari100K single game.

[0159] In step S2, the present application uses the recurrent mode of Mamba for fast inference. At this time, Mamba needs to input an additional intermediate vector x t when inputting, and output the intermediate vector x t at the next time step in addition to the recurrent state h t+1 when outputting; at this time, the present application can perform fast single-step inference like DreamerV3, a model-based algorithm based on RNN.

[0160] In step S3, the present application uses a policy gradient algorithm to train the robot arm model; unlike DreamerV3 and STORM algorithms, the present application uses Histogram encoding instead of the original TwoHot encoding when updating the value function V ψ , so that the robot arm model learns a better strategy.

[0161] Of course, the present application can have other various embodiments, and those skilled in the art can certainly make various corresponding changes and modifications according to the present application without departing from the spirit and essence of the present application. However, these corresponding changes and modifications should all belong to the protection scope of the claims attached to the present application.

Claims

1. A robotic arm manipulation method based on selective state space and model-based reinforcement learning, characterized in that, Includes the following steps: S1: Construct a world model based on the real environment. The world model includes the observation encoder q. φ (z t |o t Image decoder Motion Information Mixer m φ (z t ,a t ), sequence model f φ (e 1:T (x0) Dynamic predictor Reward Predictor Termination Predictor The world model is constructed based on the real environment as follows: S11: Real-time RGB image capture of the real environment is obtained by a fixed image sensor, yielding environmental observation data at the current time t. t ; S12: By observation encoder q φ (z t |o t Environmental observation data t The resulting random variables are represented as latent space random variables z using one-hot encoding after dimensionality reduction and random sampling. t ; S13: By image decoder For the latent space random variable z t Deconvolution is performed to obtain the reconstructed environmental observation data. S14: The latent space random variable z t And robotic arm models in environmental observation data t The following action a t splicing, a motion information mixer composed of multilayer sensing mechanisms m φ (z t ,a t The concatenated data is mapped to state e. t ; S15: Sequence Model Based on Mamba Architecture φ (e 1:T ,x0) in state e t The intermediate vector x at the previous time step t-1 As input, output the loop state h at the current time t. t and intermediate vector x t ; S16: By dynamic predictor Receive loop state h t Output the hidden space random variable at the next time step. S17: By reward predictor Receive loop state h t and latent space random variables Output the reward at the current time t. S18: Termination Predictor Receive loop state h t and latent space random variables Output the predicted value of the termination signal at the current time t. S19: Based on environmental observation data award Termination signal prediction value Construct a loss function and determine whether the loss function at the current time t is less than a set value. If yes, obtain the final world model; otherwise, proceed to step S110. S110: Adjust the observation encoder q φ (z t |o t Image decoder Motion Information Mixer m φ (z t ,a t ), sequence model f φ (e 1:T (x0) Dynamic predictor Reward Predictor Termination Predictor The parameters are then re-executed, and steps S11 to S19 are repeated until the loss function is less than the set value. S2: Interact the robotic arm model with the world model to obtain the action sequence of the robotic arm model under different environmental observations, specifically: S21: Input the historical environmental observation data o0~o7 for 8 consecutive time steps and the historical actions a0~a7 performed by the robotic arm under the environmental observation data o0~o7 into the world model to obtain the initial intermediate vector x0, the loop state h0 and the latent space random variable z1; S22: The world model passes h0 and z1 to the robotic arm model, and the robotic arm model outputs the decision action. S23: The motion information mixer will combine z1 with... The concatenation is mapped to e1. ​​The sequence model takes e1 and x0 as input and outputs the current time step's loop state h1 and intermediate vector x1. S24: Input h1 into the dynamic predictor and map h1 to a 1024-dimensional regression vector Logits; S25: Divide the regression vector Logits into 32 groups of 32-dimensional vectors and perform softmax activation function operation to output 32 groups of prior distributions; S26: Randomly sample 32 prior distributions and represent the sampling results using one-hot encoding. Then, concatenate the 32 sampling results represented by one-hot encoding back to the 1024-dimensional array to obtain the latent space random variable predicted by the dynamic predictor for the next time step. S27: h1, The predicted reward is passed to the reward predictor and then terminated to obtain the predicted reward. and termination signal prediction value S28: Repeat steps S22 to S27 until T sets of trajectory data are obtained, wherein the T sets of trajectory data include the loop state h. 0:T-1 Predicted latent space random variables Predicted rewards Predicted termination signal prediction value S3: Train the robotic arm model using motion sequences observed under different environments; S4: Use the trained robotic arm model to control the robotic arm to perform corresponding actions in different real-world environments, thereby achieving robotic arm control.

2. The robotic arm manipulation method based on selective state space and model-based reinforcement learning as described in claim 1, characterized in that, According to environmental observation data award Termination signal Constructing the loss function as follows: Where φ is the observation encoder q φ (z t |o t Image decoder Motion Information Mixer m φ (z t ,a t ), sequence model f φ (e 1:T (x0) Dynamic predictor Reward Predictor Termination Predictor The set of parameters; The observation reconstruction loss at the current time t, Predict the loss for the reward at the current time t. Predict the loss for the termination signal at the current time t. The dynamic prediction loss at the current time t, Let β1 be the representation loss at the current time t, β2 be the weight of the dynamic prediction loss, and β1 be the weight of the representation loss.

3. The robotic arm manipulation method based on selective state space and model-based reinforcement learning as described in claim 2, characterized in that, Observational reconstruction loss The calculation method is as follows: Where ||·||2 represents the calculation of Euclidean distance; Reward Prediction Loss The calculation method is as follows: Where twohot represents two-hot encoding, symlog(r t ) = sign(r t )ln(|r t |+1), Reward Logits is a reward predictor The output reward regression vector; Termination signal prediction loss The calculation method is as follows: Among them, c t This represents the actual value of the termination signal at the current time t; Dynamically predicting losses The calculation method is as follows: Characterization loss The calculation method is as follows: Where sg() is the terminating gradient operator, z t+1 Let o be the latent space random variable at the next time step. t+1 For the environmental observation data captured by the image sensor at the next moment, KL[] represents the KL divergence. Indicates that the dynamic predictor is h t For the input predicted The distribution of .

4. The robotic arm manipulation method based on selective state space and model-based reinforcement learning as described in claim 1, characterized in that, Training a robotic arm model using motion sequences observed under different environments specifically includes the following steps: S31: Extract the trajectory data from group T respectively. with h t-1 By concatenating the states, we obtain T states s. t ; S32: Rewards based on the trajectory data of group T Obtain each state s separately t The corresponding value function V ψ (s t ); S33: Separate each state s t Input a robotic arm model, output the robotic arm's decision action; S34: According to the value function V ψ (s t ), construct the training loss function of the robotic arm model based on the decision action, and determine whether the training loss function is less than the set value. If yes, the final robotic arm model is obtained; if no, proceed to step S35. S35: Change the parameters of the robotic arm model and repeat steps S31 to S34 until the training loss function is less than the set value.

5. The robotic arm manipulation method based on selective state space and model-based reinforcement learning as described in claim 4, characterized in that, The training loss function for the robotic arm model is as follows: Where θ is the parameter to be optimized in the policy function. Let ψ be the loss function corresponding to the policy function, and ψ be the parameter to be optimized in the value function. Here, is the loss function corresponding to the value function, sg() is the terminating gradient operator, B is the training batch, L is the length of the trajectory data used in each batch, and S is the set scaling factor. For state s t Select action The policy function is defined by HL() for Histogram encoding, max() for maximizing the value, and η for setting the weights. Entropy of the policy function, activation function Where V represents the value function V ψ () with s t The original output from the input is a 255-dimensional vector, v. j This represents the value of the j-th dimension of V. To estimate the true value function using time-series differencing, we have: Where λ is the set weight and γ is the attenuation coefficient.

6. A robotic arm manipulation method based on selective state space and model reinforcement learning as described in any one of claims 1 to 5, characterized in that, The movement of the robotic arm is a seven-dimensional vector. The first three dimensions represent the coordinate changes of the robotic arm's end effector along the x-axis, y-axis, and z-axis in the Cartesian coordinate system, respectively. The middle three dimensions represent the rotation angles of the robotic arm's end effector with the x-axis, y-axis, and z-axis as rotation axes, respectively. The last dimension represents the opening and closing of the robotic arm's end gripper.

Images