Training reinforcement learning agents to perform multiple tasks across diverse domains

Abstract

A method, system, and apparatus, including a computer program encoded on a computer storage medium, for training a policy neural network used to select actions to be performed by an agent interacting with an environment. In one aspect, the method includes receiving latent representations characterizing a current state of the environment, generating imagination paths for the latent representations, determining predicted rewards and generating predicted state values ​​for each latent representation in the imagination paths, determining goal state values ​​for each latent representation, determining updates to current values ​​of policy network parameters, applying a symmetric logarithmic transformation to each goal state value, encoding each transformed goal state value to generate an encoded transformed goal state value, and determining updates to current values ​​of value network parameters by optimizing a critic objective function.

Description

[Technical Field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Application No. 63 / 437,931, filed January 9, 2023. The disclosure of the prior application is considered part of the disclosure of this application and is incorporated by reference into the disclosure of this application. [Background technology]

[0002] This specification relates to reinforcement learning.

[0003] In a reinforcement learning system, an agent interacts with an environment by performing actions selected by the reinforcement learning system in response to receiving observations that characterize the current state of the environment.

[0004] Some reinforcement learning systems select an action for an agent to perform in response to receiving a given observation according to the output of a neural network.

[0005] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output for a received input. Some neural networks are deep neural networks that contain one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from the received input according to the current values ​​of its respective set of parameters. Summary of the Invention

[0006] Described herein is a system implemented as a computer program on one or more computers at one or more locations that trains a policy neural network that can be used to control an agent to interact with an environment by selecting actions for the agent to perform. The system implements a general reinforcement learning framework that can effectively train the policy neural network to select actions that enable the agent to accomplish numerous tasks across a wide range of domains.

[0007] Tasks may include, for example, navigating to a specified location in an environment, identifying a particular object in the environment, manipulating a particular object in a specified manner, etc. In some cases, a task is a long-term task that requires the agent to select and execute thousands or millions of actions in succession to complete the task. In some cases, a task is a task performed by an agent in a complex environment, e.g., a complex workspace containing multiple obstacles and / or goal objects, multiple goal locations, or both.

[0008] The domain of each task generally defines the nature and / or characteristics of various aspects of the task, e.g., the environment in which the agent is deployed, the space of possible actions the agent can perform, the data characterizing the environment (also referred to herein as "observations") received by the agent, etc. These task domains may differ from each other in one or more of these aspects, for example, in action space (e.g., continuous action space vs. discrete action space), observation space (e.g., higher dimensional visual input vs. lower dimensional text input, or 2D vs. 3D camera imagery), data budget, reward frequency, or reward scale.

[0009] In one aspect, this specification describes one or more computer-implemented reinforcement learning methods for training a policy neural network to control an agent acting in an environment based on rewards received as a result of actions taken. The policy neural network can be trained using a reinforcement learning objective function that depends on the rewards received, and the rewards, or values ​​that depend on the rewards, can be scaled or "squashed" using a symmetric logarithmic transformation as described herein.

[0010] In general, one innovative aspect of the subject matter described herein can be embodied in a method of training a policy neural network having a plurality of policy network parameters used to select an action to be taken by an agent in an environment, the method including: receiving a latent representation in a latent space characterizing a current state of the environment; and, at each of a plurality of time steps, performing the following: processing the current latent representation using the policy neural network and in accordance with current values ​​of the policy network parameters to generate an action selection output; using the action selection output to select an action that would be taken by the agent when the environment is in the state characterized by the current latent representation; and processing the current latent representation and the selected action to generate an updated latent representation characterizing a state to which the environment would transition if the agent performed the selected action when the environment is in the state characterized by the current latent representation. generating a imagination path; determining, for each latent expression of the imagination path, a predicted reward; and processing the latent expression using a value neural network having a plurality of value network parameters and according to current values ​​of the value network parameters to generate a predicted state value representing the sum of time-adjusted predicted rewards to be received by the agent by selecting an action using the action selection output starting from a state characterized by the latent expression; determining a goal state value for each latent expression of the imagination path from the predicted reward and the predicted state value for the latent expression of the imagination path; determining updates to the current values ​​of the policy network parameters by optimizing an actor objective function including a first term that depends on the goal state value; applying a symmetric logarithmic transformation to each goal state value to generate, for each latent expression of the imagination path, a transformed goal state value; and encoding each transformed goal state value to generate, for each latent expression of the imagination path, an encoded transformed goal state value.and determining updates to the current values ​​of the value network parameters by optimizing a critic objective function that assesses the difference between the predicted state values ​​and the transcoded target state values.

[0011] The symmetric logarithmic transformation can take the form sign(x)ln(|x|+1) or sign(x)log(|x|+1), where x is the target state value.

[0012] The transcoding destination state values ​​may be represented by a vector of a predetermined length having a predetermined number of non-zero entries.

[0013] Generating a predicted state value for each latent representation of the imagination pathway may include processing the latent representation to assign a respective probability to each predetermined interval in a sequence of predetermined intervals among possible ranges of state values, generating an output of a value neural network by calculating a weighted average of the state values ​​represented by the sequence of predetermined intervals weighted by their respective probabilities, and applying a symmetric exponential transformation to the output of the value neural network to generate a predicted state value for the latent representation of the imagination pathway.

[0014] The symmetric exponential transformation is sign(x)(exp(|x|)-1) or sign(x)(10 |x|-1 ), where x is the output of the value neural network.

[0015] Determining updates to the current values ​​of the policy network parameters by optimizing the actor objective function may include generating a downscaled goal state value for each latent representation of the imagination path by using the scaling factor and from the goal state value, and determining updates to the current values ​​of the policy network parameters by optimizing the actor objective function including a first term that depends on the downscaled goal state value.

[0016] Generating the downscaled target state value may include dividing the target state value by a given integer value and the maximum of a scaling factor, where the scaling factor is calculated using a percentile of each of the one or more target state values ​​for the latent representation of the imagination pathway.

[0017] The actor objective function may include a second term that depends on the scaled entropy of the action selection output produced by the policy neural network, where the scaled entropy may be scaled by a predetermined entropy scaling factor.

[0018] Processing the current latent representation and the selected action to generate an updated latent representation may include processing the current latent representation and the selected action using a sequence model neural network in accordance with current values ​​of sequence model network parameters to update a hidden state of the sequence model neural network to generate an updated hidden state that characterizes the current state of the environment, and processing the updated hidden state using a dynamics predictor neural network in accordance with current values ​​of dynamics predictor network parameters to generate an updated latent representation.

[0019] Receiving a latent representation characterizing a current state of the environment may include processing the observations characterizing the current state of the environment using an encoder neural network according to current values ​​of encoder network parameters to generate the latent representation.

[0020] The method may further include processing the latent representation and the updated hidden state using a decoder neural network according to the current values ​​of the decoder network parameters to generate a reconstruction of the observation; applying a symmetric logarithmic transformation to the observation to generate a transformed observation; and determining an update to the current values ​​of the decoder network parameters based on calculating a gradient of a prediction objective function that includes a first term that assesses the difference between the reconstruction of the observation and the transformed observation, and with respect to the decoder network parameters.

[0021] For each latent representation of the imagination pathway, determining the predicted reward may include processing the latent representation and updating hidden states of the sequence model neural network using a reward neural network according to current values ​​of reward network parameters to generate the predicted reward.

[0022] The method may further include applying a symmetric logarithmic transformation to rewards received by the agent when the environment is in a state characterized by the latent representation to generate a transformed reward, and determining an update to a current value of the reward network parameter based on calculating a gradient of a predicted objective function that includes a second term that assesses the difference between the predicted reward and the transformed reward and with respect to the reward network parameter.

[0023] The method may further include determining updates to current values ​​of the sequence model neural network parameters and the dynamics predictor neural network parameters based on calculating gradients of a dynamics objective function with respect to the sequence model neural network parameters and the dynamics predictor neural network parameters, the dynamics objective function depending on the Kullback-Leibler divergence between (i) the latent representations generated by the encoder neural network and (ii) the latent representations generated by the dynamics predictor neural network.

[0024] The encoder neural network and the decoder neural network may each be configured as a respective convolutional neural network or a multi-layer perceptron (MLP), the dynamics predictor neural network and the reward neural network may each be configured as a multi-layer perceptron (MLP), and the sequence model neural network may be configured as a recurrent neural network.

[0025] A convolutional neural network, a multilayer perceptron (MLP), or a recurrent neural network may include one or more normalization layers and one or more sigmoid linear unit (SiLU) activation layers.

[0026] The method may further include initializing current values ​​of the reward network parameters and the value network parameters to zero.

[0027] The method may further include using the policy neural network to select an action to be taken by the agent in the environment by processing inputs derived from data characterizing a current state of the environment at the time step using the policy neural network to generate an action selection output used to determine an action to be taken by the agent at the time step.

[0028] In another aspect, this specification describes a method executed by one or more computers for training an encoder-decoder neural network. The encoder-decoder neural network may include an encoder neural network configured to receive an encoder network input and generate an encoder network output, and a decoder neural network configured to receive a decoder network input including the encoder network output and generate a decoder network output. The encoder-decoder neural network is configured to receive an input data item (at the encoder network input), process the input data item to generate (an encoder network output including) a latent representation of the input data item, and process the latent representation of the input data item (using the decoder neural network) to generate (at the decoder network output) a reconstruction of the input data item. Updates to current values ​​of encoder-decoder network parameters are determined based on evaluating a loss function that assesses the difference between a reconstruction of the training data item and a transformed training data item transformed using a symmetric logarithmic transformation.

[0029] In general, other innovative aspects of the subject matter described herein can be embodied in a method executed by one or more computers for training an encoder-decoder neural network configured to receive an input data item, process the input data item to generate a latent representation of the input data item, and process the latent representation of the input data item to generate a reconstruction of the input data item, the method including receiving a training data item, processing the training data item using the encoder-decoder neural network to determine a latent representation of the training data item and generate a reconstruction of the training data item according to current values ​​of encoder-decoder network parameters, applying a symmetric logarithmic transformation to the training data item to generate a transformed training data item, and determining updates to the current values ​​of the encoder-decoder network parameters based on evaluating a loss function that assesses a difference between the reconstruction of the training data item and the transformed training data item.

[0030] The symmetric logarithmic transformation can take the form sign(x)ln(|x|+1) or sign(x)log(|x|+1), where x is a training data item.

[0031] The loss function may be a mean squared error loss function.

[0032] Generating the reconstruction of the input data item may include processing the training data item to determine a latent representation of the training data item, processing the latent representation of the training data item to generate an output of an encoder-decoder neural network, and applying a symmetric exponential transformation to the output of the encoder-decoder neural network to generate the reconstruction of the training data item.

[0033] The symmetric exponential transformation is sign(x)(exp(|x|)-1) or sign(x)(10 |x|-1 ), where x is the output of the encoder-decoder neural network.

[0034] The input data items may include images.

[0035] Other embodiments of this aspect include corresponding computer systems, devices, and computer programs stored on one or more computer storage devices, each configured to perform the actions of the method. One or more computer systems may be configured to perform particular operations or actions by software, firmware, hardware, or any combination thereof installed on the system that, when operated, causes the system to perform the actions. One or more computer programs may be configured to perform particular operations or actions by including instructions that, when executed by a data processing device, cause the device to perform the actions.

[0036] The subject matter described herein can be implemented in particular embodiments to realize one or more of the following advantages.

[0037] Many complex tasks, such as robotic tasks, involve reward scales that can vary, sometimes by large factors, depending on the application domain, the specific task, and the stage of learning. This can adversely affect stability and can result in learning delays or failure. The described techniques can implement a general reinforcement learning (RL) framework that trains a policy neural network along with one or more additional neural networks based on a latent representation of the environment to achieve robust performance across multiple tasks across a wide range of domains, including long-term tasks (tasks that require the agent to select and execute thousands or millions of actions in sequence) and even tasks performed in environments with few rewards or difficult exploration.

[0038] The described reinforcement learning framework can use a fixed set of hyperparameters that can be easily applied across these diverse domains, allowing for easy setup whenever a new task needs to be learned. Therefore, under the described framework, training an instance of a policy neural network for each new task can consume less real-time and less computational resources, e.g., memory and processing power, than traditional RL systems, because hyperparameter tuning by repeatedly running multiple training trials is no longer required. The described reinforcement learning framework is also scalable, meaning that arbitrarily large policy neural networks can be trained, and an increase in model size (e.g., in terms of model parameters) corresponds to both improved task performance and improved data efficiency.

[0039] In particular, the described framework leverages a symmetric logarithmic transformation that compresses the magnitude of both large positive and negative output values ​​of neural networks while preserving the input sign, thereby mitigating common issues, including divergence, stagnation, and non-stationarity, that can otherwise hinder RL training when some existing loss functions and regularization techniques are used; instead, it unifies gradient signals across orders of magnitude, thereby ensuring stable training of neural networks across diverse domains.

[0040] The details of one or more embodiments of the subject matter described herein are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. [Brief explanation of the drawings]

[0041] [Figure 1] 1 shows an example of a training system. [Figure 2A] 1 is an example diagram of training a world model neural network. [Figure 2B] 1 is an example diagram of training a policy neural network and a value neural network. [Figure 3] 1 shows an example of data flow for an agent interacting with an environment to select an action to perform. [Figure 4] 1 is a flow diagram of an example process for training a policy neural network and a value neural network. [Figure 5] 1 is a flow diagram of an example process for generating imagination pathways. [Figure 6] 1 is a flow diagram of an example process for training a world model neural network. [Figure 7] 1 is a flow diagram of an example process for training an encoder-decoder neural network. [Figure 8]Quantitative examples of the performance improvements that can be achieved by using the neural networks described herein are presented. DETAILED DESCRIPTION OF THE INVENTION

[0042] Like reference numbers and designations in the various drawings indicate like elements.

[0043] This specification describes a training system, implemented as a computer program on one or more computers at one or more locations, for training a policy neural network. During or after training, the policy neural network can be used to control a reinforcement learning agent that interacts with an environment by processing inputs derived from data characterizing the current state of the environment at the time step (i.e., "observations") to generate outputs that are used to determine actions the agent should take, at each of a number of time steps.

[0044] Specifically, the training system implements a general reinforcement learning (RL) framework to train a policy neural network by generating multiple imagination paths and then training the policy neural network on the multiple imagination paths, thereby training the policy neural network so that it can be used to control an agent that performs any of multiple tasks.

[0045] Each imagination path contains data characterizing a sequence of predicted state transitions of the environment caused by actions that will be performed by the agent when performing the task. To that end, each imagination path contains multiple latent representations in a latent space that characterize different states of the environment across multiple time steps in the imagination path. The multiple latent representations include a latent representation characterizing an initial state of the environment, a latent representation characterizing a final state of the environment (e.g., a goal state of the environment achieved upon completion of the task), and one or more latent representations each characterizing an intermediate state of the environment between the initial and final states.

[0046] As used herein, a "latent representation" refers to a representation of a state of an environment that is more data-efficient, e.g., more compact, than an observation of the same state of the environment. A latent representation comprises an ordered set of numerical values, e.g., a tensor of numerical values ​​(e.g., a vector or matrix), having a given dimension that is generally lower than the dimension of the observation. As used herein, the space of possible tensors having a given dimension is referred to as a "latent space."

[0047] 1 illustrates an example of a training system 100. Training system 100 is an example of a system in which the systems, components, and techniques described below are implemented as computer programs on one or more computers at one or more locations.

[0048] The training system 100 includes a policy neural network 120, a value neural network 130, and a world model neural network 140. At a given time step t in the path, the policy neural network 120 generates a latent representation z that characterizes the current state of the environment at the given time step t. t and the action a that the agent should take at a given time step t. ta neural network π configured to process policy network inputs according to a set of parameters θ (referred to herein as “policy network parameters”) to generate an action selection output that can be used to select θ The latent representation z t is later called the updated latent representation. Such an updated latent representation z t Obtaining for a given time step t is further described below.

[0049] Some examples of action selection outputs and how they can be used to select actions are described below. In one example, the action selection output may include a numerical probability value for each action in a set of possible actions that can be performed by the agent. When used to select actions, the action to be performed by the agent may be selected, for example, by sampling the actions according to their probability values ​​or by selecting the action with the highest probability value.

[0050] In other examples, the action selection output may directly define the action that the agent should perform, for example, by defining torque values ​​to be applied to the joints of the robotic agent. Such an action may be the action predicted to result in the highest long-term time-discounted reward when performed by the agent in response to the observation.

[0051] In another example, the action selection output may include a Q value for each action in a set of possible actions that can be performed by the agent. When used to select an action, the Q values ​​may be processed (e.g., using a softmax function) to generate a respective probability value for each possible action, which may be used to select an action for the agent to perform (as described above). The action with the highest Q value may also be selected as the action for the agent to perform.

[0052] In this example, the Q-value of an action is an estimate of the "return" that results from the agent taking the action in response to a current observation and then selecting future actions to be taken by the agent according to the current values ​​of the policy network parameters.

[0053] Return refers to a cumulative measure of the "reward" an agent receives, e.g., the sum of time-discounted rewards. An agent can receive a respective reward at each time step, where the reward is specified by a scalar numeric value, e.g., characterizing the agent's progress toward completing an assigned task.

[0054] At a given time step t in the path, the value neural network 130 generates an updated latent representation z that characterizes the current state of the environment at the given time step t. t and processing the value network inputs according to a set of parameters ψ (referred to herein as "value network parameters") to generate a predicted state value, i.e., a predicted value of the current state of the environment. In some implementations, the predicted state value is calculated using an updated latent representation z t represents the time-adjusted sum of predicted rewards that the agent would receive by selecting an action using the action selection output of the policy neural network during a path starting from the current state of the environment corresponding to

[0055] The world model neural network 140 can incorporate what the training system 100 knows about the environment (the "world") so that it can make predictions about paths in the latent space based on what it knows about the past. Generally, the world model neural network 140 is configured to simulate the state transitions of the environment that would be caused by actions taken by the agent, which may be selected using the action selection output of the policy neural network 120 at multiple points during the path.

[0056] World model neural network 140 may include sequence model neural network 150, an encoder-decoder neural network including encoder neural network 160 and decoder neural network 165, dynamics predictor neural network 170, reward neural network 180, and continuation predictor neural network 190. In some implementations, world model neural network 140 includes all of neural networks 150, 160, 165, 170, 180, and 190. In other implementations, world model neural network 140 includes only some of these neural networks; for example, decoder neural network 165 and continuation predictor neural network 190 may be optional components of world model neural network 140.

[0057] The sequence model neural network 150 generates a current latent representation z that characterizes the previous state of the environment at a given time step t in the path. t-1 and a previously selected action a according to a set of parameters (referred to herein as "sequence model network parameters"). t-1 and a sequence model neural network f φ The hidden state h t-1 update the previous latent representation z t-1and the previously selected action a t-1 By processing the previous latent representation z t-2 and a further previously selected action a t-2 The current state h generated by processing the data t-1 A neural network f configured to modify φ In general, the subscript φ is used herein to denote a parameter of the world model neural network 140 and any parameter of the neural network it contains. As used herein, "parameter" refers to a trainable parameter of a neural network, such as a weight, which is distinct from a hyperparameter.

[0058] At a given time step t, the hidden state of the sequence-model neural network 150 after processing the world-model network input is referred to herein as the updated hidden state h that characterizes the current state of the environment. t The current state is defined as the state where the environment is t-1 The previously chosen action a of the agent when in a previous state characterized by t-1 is the state to which the environment will transition if the sequence model neural network 150 is executed. t can be a predetermined initial hidden state.

[0059] Similar to a latent representation, a "hidden state" refers to a representation of a state of an environment that is more data-efficient, e.g., more compact, than an observation of the same state of the environment. A hidden state comprises an ordered set of numbers, e.g., a tensor of numbers (e.g., a vector or matrix), having a certain dimensionality that is generally lower than the dimensionality of the observation.

[0060] The encoder-decoder neural network includes an encoder neural network 160 that, at a given time step t, generates an updated hidden state h that characterizes the current state of the environment. tand the current observation x of the current state of the environment t and an encoder network that receives inputs and generates an updated latent representation z that characterizes the current state of the environment. t a neural network q configured to process the encoder network inputs according to a set of parameters φ (referred to herein as “encoder network parameters”) to generate φ is.

[0061] While observations of a state of the environment typically include high-dimensional sensor data, e.g., image data or video data, that characterize the state of the environment, latent representations that characterize the same state of the environment can be represented as an ordered set of numbers, e.g., a vector or matrix of numbers, that generally have lower dimensionality than the observations. For example, each observation can be an RGB image with 64x64x3 pixels, while each latent representation can be a 32-dimensional array containing 32 latent variables.

[0062] The decoder neural network 165 included in the encoder-decoder neural network generates an updated latent representation z characterizing the current state of the environment at a given time step t. t and, in some embodiments, an updated hidden state h characterizing the current state of the environment. t and a decoder network receiving inputs including a reconstructed current observation of the current state of the environment.

number

number

[0063] The dynamics predictor neural network 170 calculates, at a given time step t, an updated hidden state h that characterizes the current state of the environment. t A dynamics predictor network receives input and estimates and updates a latent representation that characterizes the current state of the environment.

number

number

[0064] Estimated updated latent representation

number

number

[0065] The reward neural network 180 generates an updated latent representation z that characterizes the current state of the environment at a given time step t. t and, in some embodiments, an updated hidden state h characterizing the current state of the environment. tand receiving a reward network input including a predicted reward.

number

number

[0066] The continuous predictor neural network 190 generates an updated latent representation z that characterizes the current state of the environment at a given time step t. t and, in some embodiments, an updated hidden state h characterizing the current state of the environment. t and a predicted path continuation flag c indicating whether the current state is a terminal state in the path, i.e., whether another updated latent representation should be generated to characterize the future state of the environment. t is a neural network p configured to process the continuous predictor network inputs according to a set of parameters φ (referred to herein as “continuous predictor network parameters”) to generate

[0067] In some embodiments, the predicted path continuation flag

number

[0068] Policy neural network 120, value neural network 130, sequence model neural network 150, encoder neural network 160, decoder neural network 165, dynamics predictor neural network 170, reward neural network 180, and continuation predictor neural network 190 may each be implemented with any suitable neural network architecture that enables them to perform the described functions.

[0069] In some implementations, policy neural network 120 and value neural network 130 may include any suitable number of layers (e.g., 5, 10, or 25 layers) of any suitable type (e.g., fully connected layers, convolutional layers, recurrent layers, attention layers, etc.) and may be connected in any suitable configuration (e.g., as a linear sequence of layers, with or without residual connections).

[0070] In some of these implementations, policy neural network 120 may be a large neural network, e.g., a neural network including a large number of layers (e.g., 50, 100, or more) and a correspondingly large number of model parameters. Generally, the larger the number of model parameters, the better the task performance and data efficiency. For example, a larger policy neural network with 18 million parameters will generally outperform a smaller policy neural network with 8 million parameters after the same number of training iterations, e.g., in terms of returns received by an agent controlled using the neural network. That is, the number of neural layers of the neural network and the number of trainable parameters of the system can be scaled according to the complexity of the task to be learned.

[0071] In some implementations, the value neural network 130 is configured to generate predicted state values ​​using a discrete regression approach. Specifically, the value neural network 130 first generates K predetermined intervals b between possible ranges of state values, e.g., using a softmax layer of the value neural network. i Probability distribution p over a sequence B of ψ (b i |s t ), where K can be any positive integer (e.g., 63, 127, 255, etc.). That is, the probability distribution p ψ (b i |s t ) is the time required for each predetermined interval b in a sequence B of K predetermined intervals. i The value neural network 130 then generates a predicted state value by calculating the average of the state values ​​represented by the sequence over a predetermined interval. The average is calculated as a weighted average, weighted by the respective probability scores.

[0072] In some implementations, sequence model neural network 150 is configured as a recurrent neural network that includes one or more recurrent layers, such as a long short-term memory (LSTM) layer or a gated recurrent unit (GRU) layer. If multiple recurrent layers are present, the hidden state may be, for example, the hidden state of the last recurrent layer in a stack of multiple recurrent layers included in sequence model neural network 150.

[0073] By way of example only, in some implementations, for example, when the observations include images, encoder neural network 160 and decoder neural network 165 may each be configured as a respective convolutional neural network (CNN) including one or more convolutional layers. In some other implementations, they may each be configured as a respective multi-layer perceptron (MLP) including one or more fully connected layers. In some implementations, one or more CNNs may be included to process one or more visual observations, and one or more MLPs may be included to process one or more lower-dimensional observations.

[0074] In some implementations, similar to encoder neural network 160 and decoder neural network 165, dynamics predictor neural network 170, reward neural network 180, and continuation predictor neural network 190 are each configured as respective multi-layer perceptrons (MLPs) including one or more fully connected layers.

[0075] In some implementations, sequence model neural network 150, encoder neural network 160, decoder neural network 165, dynamics predictor neural network 170, reward neural network 180, and continuation predictor neural network 190 may each include one or more normalization layers, e.g., one or more normalization layers, and one or more activation layers, e.g., one or more sigmoid linear unit (SiLU) activation layers.

[0076] Depending on their architecture, the above outputs can be generated in one of a variety of ways. For example, the updated hidden state h t For example, a sequence model neural network h t =f φ (h t-1 ,z t-1 ,a t-1 ) can be generated deterministically by the last recurrent layer of

[0077] As another example, the updated latent representation z t , estimated updated latent representation

number

number

number

number

number

number

number

number

number

number

number

[0078] The training system 100 includes or has access to a replay buffer 104, which stores historical paths that may be generated as a result of an agent's (or other agents') interactions with an environment (or other instances of an environment) in an attempt to perform one or more tasks. The replay buffer 104 may be any available memory, such as one or more physical or logical storage devices.

[0079] In some implementations, each history path may include a tuple at each of multiple time points between t=1 and t=T in the history path. At a given time t, a tuple may include: (i) a current observation x that characterizes the current state of the environment; t and (ii) the current action performed by the agent, a t and (iii) the current reward received by the agent, r t and (iv) a ground truth continuation flag c that indicates whether the current state of the environment is the terminal state of the environment (e.g., the goal state of the environment that will be achieved upon completion of the task). t Includes:

[0080] Training engine 110 of training system 100 repeatedly selects history paths from replay buffer 140 and trains world model neural network 140 based on the selected history paths using appropriate world model learning techniques, described below. After training world model neural network 140, training engine 110 uses the (trained) world model neural network 140 to train policy neural network 120 and value neural network 130. This is done by using the (trained) world model neural network 140 to generate multiple imaginary paths, and then jointly training policy neural network 120 and value neural network 130 based on the multiple imaginary paths using appropriate reinforcement learning techniques.

[0081] Training policy neural network 120 and value neural network 130 based on multiple imagination pathways means that policy neural network 120 and value neural network 130 are each configured to receive latent representations as input during training (e.g., rather than directly processing actual observations, e.g., sensor data, e.g., image or video data, that characterize environmental conditions).

[0082] An embodiment of the training system 100 includes a symmetric logarithmic transform engine 112 that applies a symmetric logarithmic transform to an input x. In general, a symmetric logarithmic transform can be characterized as a logarithmic transform, i.e., a transform that applies a logarithmic function to its input or argument x to produce a transformed output, and that is symmetric about the origin (the zero value of the input and the log-transformed input) while preserving the sign of the input or argument x. Logarithms can have any base. As an example, a symmetric logarithmic transform can take the form: sign(x)ln(|x|+1) where x is an input or argument. As an example, a symmetric logarithmic transformation may use base 10 (or other base) logarithms, and may take the form, for example: sign(x)log(|x|+1)

[0083] The symmetric logarithmic transformation compresses the magnitude of both large positive and negative values ​​while preserving the input sign. By leveraging such a symmetric logarithmic transformation, training engine 110 mitigates common problems, including divergence, stagnation, and non-stationarity, that can otherwise hinder training when some existing loss functions and regularization techniques are used, and instead unifies gradient signals across orders of magnitude, ensuring stable training of neural networks across diverse domains.

[0084] 2A is an example diagram of training a world model neural network. For example, the world model neural network may correspond to world model neural network 140 of FIG. 1, which may include sequence model neural network 150, an encoder-decoder neural network including encoder neural network 160 and decoder neural network 165, dynamics predictor neural network 170, reward neural network 180, and continuation predictor neural network 190.

[0085] The training engine 110 performs the operations shown in FIG. 2A for each of a plurality of history paths selected from the replay buffer 104 to train these neural networks using appropriate world model learning techniques to enable these neural networks to more accurately simulate the environmental effects of actions performed by the agent when interacting with the environment.

[0086] As shown, at each of multiple time steps t in a given history path selected from the replay buffer 104, the sequence model neural network generates a current latent representation z characterizing the previous state of the environment. t-1 and the previously selected action a t-1 and a sequence model network that receives inputs and updates the hidden state h that characterizes the current state of the environment. t The sequence model network processes inputs to generate a current state of the environment, which is the state the environment would transition to if the agent performed the previously selected action, given that the environment was in the previous state characterized by the current latent representation.

[0087] For example, at time step t=2, the sequence model neural network generates an updated hidden state h2=f that characterizes the current state of the environment. φThe current state of the environment at time step t=2 is the state the environment would transition to if the agent performed the previously selected action a1, when the environment was in the previous state characterized by the current latent representation z1.

[0088] At each of multiple time steps t in a given history path selected from the replay buffer 104, the encoder neural network (“enc”) generates an updated hidden state h that characterizes the current state of the environment. t and the current observation x of the current state of the environment t and an encoder network that receives input and generates an updated latent representation z that characterizes the current state of the environment. t ~q φ (z t |h t ,x t )

[0089] For example, at time step t=2, the encoder neural network ("enc") generates an updated latent representation z2~q that characterizes the current state of the environment. φ The current state of the environment at time step t=2 is the state to which the environment would transition if the agent performed the previously selected action a1, given that the environment was in the previous state characterized by the current latent representation z1.

[0090] At each of multiple time steps t in a given history path selected from the replay buffer 104, a decoder neural network ("dec") generates an updated latent representation z characterizing the current state of the environment. t and, in some embodiments, an updated hidden state h characterizing the current state of the environment. t and a decoder network that receives inputs and reconstructs the current observations of the current state of the environment.

number

number

[0091] For example, at time step t=2, the decoder neural network ("dec") receives the reconstructed current observation

number

number

[0092] At each of multiple time steps t in a given history path selected from the replay buffer 104, the dynamics predictor neural network generates an updated hidden state h that characterizes the current state of the environment. t A dynamics predictor network that receives input and estimates and updates latent representations that characterize the current state of the environment.

number

number

[0093] For example, at time step t=2, the dynamics predictor neural network generates an estimated updated latent representation that characterizes the current state of the environment.

number

number

[0094] At each of multiple time steps t in a given history path selected from the replay buffer 104, the reward neural network generates an updated latent representation z characterizing the current state of the environment. t and, in some embodiments, an updated hidden state h characterizing the current state of the environment. t and a reward network input including a predicted reward

number

number

number

[0095] At each of a number of time steps t in a given history path selected from the replay buffer 104, the continuation predictor neural network generates an updated latent representation z characterizing the current state of the environment. t and, in some embodiments, an updated hidden state h characterizing the current state of the environment. tand a continuation predictor network input including a predicted path continuation flag indicating whether the current state is a terminal state in the path.

number

[0096] For example, at time step t=2, the continuation predictor neural network sets a predicted path continuation flag indicating that the current state is not the terminal state of the history path.

number

[0097] For a given historical path selected from the replay buffer 104, the training engine 110 calculates a predicted objective function L pred Based on optimizing (φ), the values ​​of the parameters φ of the dynamics predictor neural network, the reward neural network, and the continuation predictor neural network included in the world model neural network are updated.

[0098] Prediction objective function L pred (φ) is, for each of several time steps t in a given path, (i) the reconstructed current observation

number

[0099] In particular, to generate the transformed current observation, the simlog transformation engine 112 converts the current observation x t Apply a symmetric logarithmic transformation to . For example, the symmetric logarithmic transformation can take the form: sign(x t )ln(|x t |+1) or sign(x t )log(|x t |+1) In the formula, x t is the current observation. The symmetric logarithmic transformation can be applied to any type of observation, including, for example, low-dimensional inputs. When applied to a vector or matrix, the symmetric logarithmic transformation can be applied to each element of the vector or matrix.

[0100] The first term trains the decoder neural network to generate a reconstructed current observation that is closest to the current observation contained in the history path, trains the encoder neural network to generate updated semantically significant latent representations by backpropagation, and trains the sequence model neural network to generate updated semantically significant hidden states.

[0101] Prediction objective function L pred (φ) is the predicted reward for each of the time steps t in a given path.

number

[0102] Prediction objective function L pred (φ) is the predicted path continuation flag c for each of multiple time steps t in a given path. t and the ground truth path continuation flag, which trains the continuation predictor neural network to generate a more accurate predicted path continuation, i.e., to generate a predicted path continuation flag that is closest to the path continuation flag contained in the history path, trains the encoder neural network to generate, by backpropagation, semantically significant updated latent representations, and trains the sequence model neural network to generate semantically significant updated hidden states.

[0103] For example, the predicted objective function L pred (φ) can be expressed as follows:

number

[0104] For a given path selected from the replay buffer 104, the training engine 110 also calculates the dynamics objective function L dynBased on optimizing (φ), the values ​​of the parameters φ of the dynamics predictor neural network included in the world model neural network 140 are updated.

[0105] Dynamics objective function L dyn (φ) is the updated latent representation z generated by the encoder neural network for each of the time steps t in a given path. t and (ii) the estimated updated latent representation generated by the dynamics predictor neural network.

number

[0106] For example, the difference can be determined as the Kullback-Leibler (KL) divergence, and the dynamics objective function L dyn (φ) can be:

number

[0107] For a given path selected from the playback buffer 104, the training engine 110 calculates the representation objective function L repBased on optimizing (φ), the value of the parameter φ of the encoder neural network included in the world model neural network is updated.

[0108] Dynamics objective function L dyn Similarly to (φ), the representation objective function L rep (φ) is the updated latent representation z generated by the encoder neural network for each of the time steps t in a given path. t and (ii) the estimated updated latent representation generated by the dynamics predictor neural network.

number

[0109] For example, the difference can be determined as the Kullback-Leibler (KL) divergence, expressing the objective function L rep (φ) can be:

number

[0110] In some implementations, the training engine 110 trains the world model neural network end-to-end, i.e., the complete world model neural network 140 trains the world model neural network 140 using the above prediction objective function L pred (φ), the dynamics objective function L dyn (φ), and the representation objective function L rep (φ) is trained by optimizing a composite objective function that combines the objective functions. For example, the composite objective function can be:

number

[0111] In this example, β pred , β dyn , and β rep are the prediction objective functions L pred (φ), the dynamics objective function L dyn (φ), and the representation objective function L rep (φ). In the embodiment, β pred , β dyn , and β rep are the hyperparameters of the training system 100 (some example hyperparameter values ​​are given later).

[0112] 2B is an example diagram of training the policy neural network and the value neural network. For example, the policy neural network and the value neural network may correspond to the policy neural network 120 and the value neural network 130, respectively, of FIG. 1. After training the world model neural network 140, the training engine 110 trains the policy neural network 120 and the value neural network 130 on the imagination paths generated using the world model neural network 140.

[0113] The training engine 110 performs the operations shown in FIG. 2B for each of the multiple imagination paths by using a world model neural network in conjunction with a policy neural network and a value neural network using appropriate reinforcement learning techniques that enable the policy neural network to generate action selection outputs that can be used to select actions that will cause the agent to effectively perform one or more tasks.

[0114] In some implementations, each imagination path may include a tuple at each of multiple time points between t=1 and t=T in the imagination path. At a given time point t, the tuple includes (i) an updated latent representation z that characterizes the current state of the environment; t and (ii) an updated hidden state h that characterizes the current state of the environment. t and (iii) the current action a selected by the action selection output. t and (iii) predicted current reward r t and (iv) a prediction continuation flag c indicating whether the current state of the environment is the terminal state of the environment (e.g., the goal state of the environment to be achieved upon completion of the task). t Includes:

[0115] As shown in the figure, the imagination path updates the latent representation z t At the first time step t=1, the updated latent representation z1~q φ (z1|h1,x1) is generated by the encoder neural network based on a given initial hidden state h1 of the sequence model neural network and the current observation x1 contained in the history path that characterizes the initial state of the environment. At any subsequent time step, the updated latent representation is generated by the updated hidden state h1 of the sequence model neural network. t Based on the estimated updated latent representation generated by the dynamics predictor neural network,

number

[0116] Thus, imagination paths differ from history paths stored in playback buffer 104 in that each imagination path includes at least one updated latent representation that is not generated based on observations of environmental conditions, e.g., sensor data, e.g., image or video data. Rather, imagination paths include updated latent representations that are generated based on hidden states of a sequence model neural network.

[0117] In general, training engine 110 uses an actor-critic technique to train a policy neural network and a value neural network, where the policy neural network corresponds to the actor neural network and the value neural network corresponds to the critic neural network.

[0118] At each of multiple time steps t in a given imagination path, the policy neural network generates an updated latent representation z that characterizes the current state of the environment at time step t. t receives a policy network input containing the action a to be taken by the agent at time step t t and the action selection output π θ (a t |s t )(where s t ={h t ,z t For example, at time step t=2, the policy neural network processes the policy network input to generate an action selection output π that can be used to select action a2 at time step t=2. θ Generate (a2|s2).

[0119] At each of multiple time steps t in a given imagination path, the value neural network generates an updated latent representation z that characterizes the current state of the environment at time step t. t receives a value network input containing a predicted state value v ψ (s t ) at time step t=2, the value neural network processes the value network inputs to generate a predicted value v of the current state of the environment at time step t=2. ψ (s t ) (referred to as "v2" in Figure 2B).

[0120] The value neural network determines K predetermined intervals b between the possible ranges of state values. i Probability distribution p over a sequence B of ψ (b i |s t In implementations that produce a symmetric logarithmic transformation ( ) as an output, the training engine 110 can use a symmetric exponential transformation engine 112 to apply a symmetric exponential transformation to the output of the value neural network to generate predicted state values. In general, the symmetric exponential transformation is the inverse of the symmetric logarithmic transformation described above.

[0121] For example, the symmetric exponential transformation symexp(x) is either sign(x)(exp(|x|)-1) or sign(x)(10 |x|-1 ) where x is the output of the Value Neural Network. The predicted state value for each of multiple time steps t in a given imagination path can be calculated as follows:

number

[0122] For a given imagination path, the training engine 110 determines a target state value at each time step t and updates the value of the parameter θ of the policy neural network based on optimizing an actor objective function. The actor objective function includes a first term that depends on the target state value. The actor objective function trains the policy neural network to produce action selection outputs that may facilitate the selection of more effective actions, which may result in higher predicted state values ​​calculated by the value neural network and higher predicted rewards calculated by the reward neural network.

[0123] The goal state value at each time step t can be calculated using the predicted reward generated by the reward neural network, the predicted state value generated by using the value neural network, or the predicted reward

number

number

[0124] In some implementations, the first term in the actor objective function relies on downscaled goal state values ​​to better accommodate sparse rewards. Specifically, the training engine 110 scales down larger goal state values ​​without scaling up smaller goal state values. This is done by dividing each goal state value by (i) a given integer value (e.g., 1, 2, 10, etc.) and (ii) the maximum value of the scaling factor S.

[0125] For example, the scaling factor S can be calculated using the respective quantiles, e.g., percentiles, of one or more goal state values ​​for the latent expressions in the path. As a specific example, the scaling factor S can be calculated as an exponentially decaying average of the range from the 5th to the 95th batch percentile.

number

[0126] In some implementations, the actor objective function is the action selection output π generated by the policy neural network. θ (a t |s t ) includes a second entropy normalization term that depends on the entropy H of the vector σ. The entropy H may be a scaled entropy that is scaled by a predetermined entropy scaling factor η.

[0127] For example, an actor's objective function might be:

number

[0128] For a given imagination path, the training engine 110 updates the values ​​of the parameters of the value neural network based on optimizing a critic objective function, which is: (i) the predicted state value v ψ (s t ) and (ii) the encoded target state value generated by the sim-log transform engine 112 from the target state value. The critic objective function trains the value neural network to generate more accurate predicted state values, i.e., predicted state values ​​that more accurately represent the time-adjusted total predicted reward that the agent will receive during a given imagination path.

[0129] In particular, to generate the encoded transformed target state values, the training engine 110 uses a symmetric logarithmic transform engine 112 that applies a symmetric logarithmic transform to the target state values ​​for each time step in the imagination path to generate transformed target state values ​​for the time step in the imagination path, and encodes the transformed target state values ​​using an appropriate encoding technique to generate encoded transformed target state values ​​for the time step in the imagination path.

[0130] As explained above, the symmetric logarithmic transformation can take the form, for example, sign(x)ln(|x|+1) or sign(x)log(|x|+1), where x is the target state value.

[0131] In an embodiment, the transcoding target state value may be represented by a vector of a predetermined length with a predetermined number of non-zero entries. For example, the training engine may use a two-hot encoding technique. The two-hot encoding technique allows the value neural network to maintain a fixed distribution over potential returns, avoiding a wide range of return distributions that could otherwise slow training. In this example, the transcoding target state value is represented by a vector of length |B|, with all elements being zero except for the two entries at positions k and k+1 that are closest to the encoded consecutive number. These two entries sum to 1, giving more weight to the entry closer to the encoded number.

number

[0132] For example, the critic objective function might be:

number

[0133] 1 , to aid in training policy neural network 120, value neural network 130, and world model neural network 140, training system 100 maintains predetermined values ​​for a set of hyperparameters 118, for example, in one or more physical or logical storage devices accessible by the training system. Hyperparameters are parameters of a neural network that are set before the neural network is trained. Hyperparameters can affect the quality of neural network training, and therefore, setting optimal values ​​for hyperparameters can be an important component of training.

[0134] In some implementations, these predetermined values ​​may be adjusted before training begins, for example, through an automated search process or manual hyperparameter fine-tuning, while in other implementations, these predetermined values ​​may be fixed and training system 100 repeatedly applies a set of hyperparameters having the same predetermined values ​​across different tasks for which the neural network is trained. In these other implementations, the predetermined values ​​of the set of hyperparameters 118 maintained by training system 100 alleviates the need for hyperparameter fine-tuning, which can be time-consuming and computationally expensive, and enables neural networks to be trained more efficiently across a range of different task domains.

[0135] Table 1 below lists an example set of hyperparameters and their predetermined values. In the table, "number of latents" refers to the dimensionality of each latent representation (which is a 32-dimensional array specifying the probability distribution of each latent out of a set of 32 possible latent variables), "horizon" refers to the number of time steps included in each imagination path, and EMA stands for Exponential Moving Average. [Table 1]

[0136] By performing the operations described in Figures 2A-2B, training engine 110 may identify trained values ​​for parameters 114 of the neural networks included in training system 100, i.e., parameters of policy neural network 120, value neural network 130, and world model neural network 140, which in turn may include sequence model neural network 150, encoder neural network 160, decoder neural network 165, dynamics predictor neural network 170, reward neural network 180, and continuation predictor neural network 190.

[0137] Once trained, some or all of the neural networks included in training system 100 can be deployed and used to control an agent to perform one or more tasks by interacting with the environment. For example, sequence model neural network 150 and encoder neural network 160 can work together to generate updated latent representations for each new observation. For example, sequence model neural network 150 can update the hidden state of the sequence model neural network to generate an updated hidden state h t The encoder neural network 160 may generate a latent representation z tTo generate the t The policy neural network 120 may then process the observation x to generate an action selection output that is used to select an action to be taken by the agent in response to a new observation. t The latent representation of z t can be processed.

[0138] In some implementations, the training system 100 may provide, e.g., via a wired or wireless connection, the trained values ​​of some or all of the network parameters 114 to an external system, which can cause the agent to perform different actions, e.g., by sending control signals directly to the agent or by sending data identifying the actions to the agent's control system.

[0139] FIG. 3 illustrates an exemplary data flow for using a policy neural network 320, a sequence model neural network 350, and an encoder neural network 360 to select an action 308 to be taken by an agent 302 interacting with an environment 304 at each of multiple time steps.

[0140] Policy neural network 320, sequence model neural network 350, and encoder neural network 360 may correspond to policy neural network 120, sequence model neural network 150, and encoder neural network 160, respectively, of Figure 1. Trained values ​​for parameters 114 may be determined (and output) by training system 100 of Figure 1.

[0141] At each time step, the sequence model neural network 350 processes the sequence model network input, including the current observation 306, according to the trained values ​​of the sequence model neural network's parameters to generate an updated hidden state 352 that characterizes the current state of the environment 304.

[0142] The encoder neural network 360 processes the encoder network inputs, including the updated hidden state 352 and the current observation 306, according to the trained values ​​of the parameters of the encoder neural network to generate an updated latent representation 362 that characterizes the current state of the environment 304.

[0143] The policy neural network 320 processes the policy network input, including the updated latent representation 362, according to the trained values ​​of the policy neural network's parameters to generate an action selection output 322 that can be used to select an action 308 to be taken by the agent 302 in response to the current observation 306. Examples of action selection outputs and how the action selection outputs are used to select an action are described above with reference to FIG.

[0144] Specifically, during inference, each received observation has a one-to-one correspondence with an updated hidden state or updated latent representation, as shown in Example Figure 3. Note that this differs from training via latent imagination, where each observation typically has a one-to-many correspondence with an updated hidden state or updated latent representation, i.e., a sequence of multiple updated hidden states or updated latent representations is generated from a single observation during training.

[0145] In some implementations, the environment is a real-world environment, the agent is a mechanical agent (or electromechanical agent) that interacts with the real-world environment, such as a robot or an autonomous or semi-autonomous vehicle for land, air, or sea that operates or navigates within the environment, and the action is an action performed by the mechanical agent within the real-world environment to perform a task. For example, the agent may be a robot that interacts with the environment to accomplish a particular task, such as locating or manipulating an object of interest within the environment, or moving an object of interest to a particular position within the environment, or navigating to a particular destination within the environment.

[0146] In these implementations, the observations may include, for example, one or more of images, object position data, and sensor data, such as sensor data from image, distance, or position sensors, or from actuators, to capture observations as the agent interacts with the environment. For example, in the case of a robot, the observations may include data characterizing the current state of the robot, such as one or more of joint positions, joint velocities, joint forces, torques or accelerations, e.g., torque feedback for gravity compensation, and the global or relative pose of an item held by the robot. In the case of a robot or other mechanical agent or vehicle, the observations may similarly include one or more of positions, linear or angular velocities, forces, torques or accelerations, and the global or relative pose of one or more parts of the agent. The observations may be defined in one, two, or three dimensions and may be absolute and / or relative observations. The observations may also include sensed electronic signals, such as, for example, motor current or temperature signals, and / or image or video data captured, for example, by a camera or LIDAR sensor, e.g., data from sensors on the agent or from sensors positioned separately from the agent in the environment.

[0147] In these implementations, an action may be a control signal or higher-level control command for controlling a robot or other mechanical agent, such as torque for a robot joint, or for controlling an autonomous or semi-autonomous vehicle for land, air, or sea, such as torque for a control surface or other control element, such as a steering control element of the vehicle. A control signal may include, for example, position, velocity, or force / torque / acceleration data for one or more joints of a robot or part of another mechanical agent. A control signal may also or alternatively include electronic control data, such as motor control data or, more generally, data for controlling one or more electronic devices in an environment, the control of which affects observed conditions of the environment. For example, in the case of an autonomous or semi-autonomous vehicle for land, air, or sea, a control signal may define actions for controlling navigation, such as steering, and movement, such as braking and / or accelerating, for the vehicle.

[0148] In some implementations, the environment is a simulation of the aforementioned real-world environment, and the agent is implemented as one or more computers that interact with the simulated environment. For example, the simulated environment may be a simulation of a robot or vehicle, and the reinforcement learning system may be trained in the simulation and, once trained, may then be used in the real world.

[0149] In some embodiments, the environment is a real-world manufacturing environment for producing a product, such as a chemical, biological, or mechanical product, or a food product. As used herein, "manufacturing" a product also includes purifying starting materials to make the product or processing starting materials, e.g., to remove contaminants or to produce a cleaned or regenerated product. A manufacturing plant may include multiple manufacturing units, such as vessels for chemical or biological substances or machines, e.g., robots, for processing solid or other materials. The manufacturing units are configured such that intermediate versions or components of the product can be moved between manufacturing units during the production of the product, e.g., via pipes or mechanical transport. As used herein, manufacturing a product also includes the production of food by a cooking robot.

[0150] The agent may include an electronic agent configured to control a manufacturing unit or a machine, such as a robot, that operates to manufacture the product. That is, the agent may comprise a control system configured to control the manufacture of a chemical, biological, or mechanical product. For example, the control system may be configured to control one or more of the manufacturing units or machines, or to control the movement of intermediate versions or components of the product between the manufacturing units or machines.

[0151] As one example, the tasks performed by the agents may include tasks to manufacture a product or an intermediate version or component of a product. As another example, the tasks performed by the agents may include tasks to control, e.g., minimize, resource usage, such as tasks to control electricity consumption, or water consumption, or consumption of any materials or consumables used in the manufacturing process.

[0152] Actions may include control actions for controlling the use of a machine or manufacturing unit that processes solid or liquid materials to produce a product or a product intermediate or component, or for controlling the movement of a product intermediate or component within a manufacturing environment, e.g., between manufacturing units or machines. Typically, actions may be any action that affects the observed state of the environment, for example, an action configured to adjust any of the sensed parameters described below. These may include actions for adjusting the physical or chemical conditions of a manufacturing unit or for controlling the movement of a mechanical part of a machine or a joint of a robot. Actions may include actions that impose operating conditions on a manufacturing unit or machine, or actions that effect a setting change to adjust, control, or switch on or off the operation of a manufacturing unit or machine.

[0153] In some implementations, the reinforcement learning context includes a reward computation unit that generates a reward (e.g., a reward in the form of a number), typically from the observation. The reward may be related to a metric of task performance. For example, for a task to manufacture a product, the metric may include the quantity of product manufactured, the quality of the product, the production rate of the product, or a metric of the physical cost of performing the manufacturing task, e.g., the amount of energy, material, or other resources used to perform the task. For a task to control resource usage, the metric may include any metric of resource usage. For a task to control an electromechanical agent, such as a robot, to perform manipulation of an object, the reward may indicate whether the object was correctly manipulated according to predefined criteria.

[0154] Typically, observations of the state of the environment can include any electronic signals representative of the function of electronic and / or mechanical items of equipment. For example, a representation of the state of the environment can be derived from observations made by sensors that sense the state of the manufacturing environment, such as sensors that sense the state or configuration of a manufacturing unit or machine, or sensors that sense the movement of materials between manufacturing units or machines. As some examples, such sensors can be sensors that can be configured to sense mechanical movement or force, pressure, temperature; electrical conditions such as current, voltage, frequency, impedance, etc.; quantity, level, flow rate / movement rate, or flow path / movement path of one or more materials; physical or chemical conditions, such as the physical state, shape, or configuration, or chemical conditions such as pH; mechanical configuration of a unit or machine, or configuration of a unit or machine, such as a valve configuration; image or video sensors that capture image or video observations of a manufacturing unit or machine or movement; or any other suitable type of sensor. If the agent is a machine such as a robot, the observations from the sensors may include observations regarding the position, linear or angular velocity, force, torque or acceleration, or pose of one or more parts of the machine, e.g., data characterizing the current state of the machine or robot, or an item held or handled by the machine or robot. The observations may also include, for example, sensed electronic signals such as motor current or temperature signals, or image data or video data, for example, from a camera or LIDAR sensor (e.g., mounted on the machine). Such sensors may be part of the agent in the environment or may be located separately from the agent.

[0155] In some implementations, the environment is a real-world environment of a service facility with multiple items of electronic equipment, such as a server farm or data center, e.g., a communications data center, or a computer data center for storing or processing data, or any service facility. The service facility may also include auxiliary controls that control the operating environment of the items of equipment, e.g., temperature controls such as refrigeration equipment, or environmental controls such as airflow controls or air conditioning equipment. The tasks may include tasks that control, e.g., minimize, resource usage, such as controlling power consumption or water consumption. The agents may include electronic agents configured to control the operation of the items of equipment or to control the operation of auxiliary controls, e.g., environmental controls.

[0156] In general, the actions may be any actions that affect the observed state of the environment, for example, actions configured to adjust any of the sensed parameters described below. These actions may include actions that control an item of equipment or an auxiliary control, or actions that impose operating conditions on an item of equipment or an auxiliary control, for example, actions that effect a setting change to adjust, control, or switch on or off the operation of an item of equipment or an auxiliary control.

[0157] Typically, observations of environmental conditions may include any electronic signals representative of the functioning of the facility or of equipment within the facility. For example, a representation of the environmental condition may be derived from observations made by any sensors that sense the state of the facility's physical environment, or by any sensors that sense the state of one or more items of equipment or one or more items of auxiliary control equipment. These sensors include sensors configured to sense electrical conditions such as current, voltage, power, or energy; facility temperature; fluid flow, temperature, or pressure within the facility or in the facility's cooling system; or physical facility configuration, such as whether a vent is open.

[0158] The reward may be related to a metric of the task's performance related to the efficient operation of the facility. For example, in the case of a task that controls, e.g., minimizes, resource usage, such as a task that controls power or water usage, the metric may include any metric related to resource usage.

[0159] In some implementations, the environment is a real-world environment of a power generation facility, e.g., a renewable power generation facility such as a solar or wind power plant. The tasks may include a control task of controlling power generated by the facility, e.g., controlling the delivery of power to a power grid, e.g., meeting grid demand or reducing the risk of mismatches between grid elements, or maximizing power generated by the facility. The agents may include electronic agents configured to control the generation of power by the facility or the coupling of generated power to a power grid. The actions may include actions of controlling the electrical or mechanical configuration of a generator, such as the electrical or mechanical configuration of one or more renewable power generation elements, e.g., controlling the configuration of a wind turbine or solar panel(s) or solar mirror, or the electrical or mechanical configuration of a rotary power generation machine. Mechanical control actions may include, e.g., actions of controlling the conversion of an energy input to an electrical energy output, e.g., controlling the conversion efficiency or degree of coupling of an energy input to an electrical energy output. Electrical control actions may include, e.g., actions of controlling one or more of the voltage, current, frequency, or phase of the generated power.

[0160] The reward may be related to a metric of performance of the task related to power distribution. For example, for the task of controlling the delivery of power to the power grid, the metric may be related to a measure of power transferred, or a measure of an electrical mismatch between a power generation facility and the power grid, such as a voltage, current, frequency, or phase mismatch, or a measure of power or energy loss at the power generation facility. For the task of maximizing the delivery of power to the power grid, the metric may be related to a measure of power or energy transferred to the power grid, or a measure of power or energy loss at the power generation facility.

[0161] Typically, observations of environmental conditions may include any electronic signals representative of the electrical or mechanical function of power generation equipment within the power generation facility. For example, a representation of the environmental condition may be derived from observations made by any sensors that sense the physical or electrical condition of equipment within the power generation facility that is generating electrical power, or the physical environment of such equipment, or the condition of auxiliary equipment supporting the power generation equipment. Thus, such observations may include observations of wind levels or solar irradiance, or observations of local time, date, or season. Such sensors may include sensors configured to sense electrical conditions of the equipment, such as current, voltage, power, or energy; temperature or cooling of the physical environment; fluid flow; or the physical configuration of the equipment; and observations of electrical conditions of the electrical grid, for example, from local or remote sensors. Observations of environmental conditions may also include one or more predictions regarding future operating conditions of the power generation equipment, such as predictions of future wind levels or solar irradiance, or predictions of future electrical conditions of the electrical grid.

[0162] In another example, the environment may be a chemical synthesis environment or a protein folding environment, whereby the states are states of a protein chain or states of one or more intermediate or precursor chemicals, and the agent is a computer system for determining how to fold a protein chain or synthesize a chemical. In this example, the actions are possible folding actions for folding a protein chain or actions for constructing precursor chemicals / intermediates, and the results to be achieved may include, for example, folding a protein to stabilize the protein and achieve a specific biological function, or providing an effective synthetic route for a chemical. In another example, the agent may be a machine agent that indirectly performs or controls protein folding actions or chemical synthesis steps, for example, by automatically controlling system-selected synthesis steps without human interaction. The observations may include direct or indirect observation of the state of the protein or chemical / intermediate / precursor and / or may be derived from simulation. Thus, the system may be used to automatically synthesize proteins with specific functional groups, such as ligands with binding sites shaped with sufficient affinity for a biological effect that can be used as a drug. For example, it may be a receptor or enzyme agonist or antagonist, or it may be an antibody configured to bind to an antibody target such as a viral coat protein or a protein expressed on cancer cells, e.g., to act as an agonist of a particular receptor or to prevent binding of other ligands and thus activation of the associated biological pathway.

[0163] Similarly, the environment can be a drug design environment, whereby the states are states of a potential pharmaceutically active compound, and the agent is a computer system that identifies components of the pharmaceutically active compound and / or synthetic pathways of the pharmaceutically active compound. The drug / synthesis can be designed based on rewards derived from targets of the pharmaceutically active compound, for example, in a simulation. As another example, the agent can be a machine agent that performs or controls the synthesis of a pharmaceutically active compound.

[0164] In some further applications, the environment is a real-world environment and the agent manages task distribution across computing resources, for example, on mobile devices and / or in a data center. In these implementations, the action may include assigning the task to a particular computing resource. In these applications, the observation may include observing computing resources such as computational and / or memory capacity, or internet-accessible resources, and the action may include assigning the task to a particular computing resource. The reward(s) may be configured to maximize or minimize one or more of computing resource utilization, power, bandwidth, and computation speed.

[0165] As a further example, the action may include presenting an advertisement, the observation may include an impression or click-through count or rate of the advertisement, and the reward may characterize a pre-selection of an item or content obtained by one or more users.

[0166] In some cases, the observations may include text or voice commands given to the agent by a third party (e.g., an operator of the agent). For example, the agent may be an autonomous vehicle, and a user of the autonomous vehicle may give the agent text or voice commands (e.g., commands to navigate to a particular location).

[0167] As another example, the environment may be an electrical, mechanical, or electromechanical design environment, e.g., an environment in which the design of an electrical, mechanical, or electromechanical entity is simulated. The simulated environment may be a simulation of a real-world environment in which the entity is intended to operate. The task may be to design the entity. The observations may include observations characterizing the entity, i.e., observations of the entity's mechanical shape, or observations of its electrical, mechanical, or electromechanical configuration, or observations of the entity's parameters or characteristics. The actions may include actions that modify the entity, e.g., actions that modify one or more of the observations. The reward may include one or more metrics related to the performance of the entity's design. For example, the reward may be related to one or more physical characteristics of the entity, such as weight or strength, or one or more electrical characteristics of the entity, such as a measure of its efficiency in performing a particular function for which the entity is designed. The design process may include outputting a design for manufacturing, e.g., in the form of computer-executable instructions for manufacturing the entity. The process may include fabricating the entity according to the design. Thus, the design of the entity may be optimized, for example, by reinforcement learning, and then the optimized design may be output, for example, as computer-executable instructions, to manufacture the entity, and then the entity may be manufactured with the optimized design.

[0168] As mentioned above, the environment may be a simulated environment. Typically, in a simulated environment, the observations may include simulated versions of one or more of the aforementioned observations or various types of observations, and the actions may include simulated versions of one or more of the aforementioned actions or various types of actions. For example, the simulated environment may be a motion simulation environment, such as a driving simulation or a flight simulation, and the agent may be a simulated vehicle that navigates through the motion simulation. In these implementations, the actions may be control inputs for controlling a simulated user or a simulated vehicle. Typically, the agent may be implemented as one or more computers that interact with the simulated environment.

[0169] The simulated environment may be a simulation of a particular real-world environment and agent. For example, the system may be used to select actions within the simulated environment during training or evaluation of the system, and after training or evaluation, or both, are complete, and an action selection policy may be deployed to control a real-world agent in the particular real-world environment that was the subject of the simulation. This may avoid unnecessary wear and tear on the real-world environment or the real-world agent, and may enable training and evaluation of a control neural network in situations that rarely occur or are difficult or unsafe to replicate in the real-world environment. For example, the system may be partially trained using a simulation of a machine agent in a simulation of a particular real-world environment, and then deployed to control an actual machine agent in the particular real-world environment. Thus, in such cases, observations of the simulated environment relate to the real-world environment, and actions selected in the simulated environment relate to actions to be performed by the machine agent in the real-world environment.

[0170] In some implementations, as mentioned above, the agent may not include a human (e.g., the agent is a robot). Conversely, in some implementations, the agent includes a human user of a digital assistant, such as a smart speaker, smart display, or other device. Information defining the task can then be obtained from the digital assistant, and the digital assistant can be used to instruct the user based on the task.

[0171] For example, a reinforcement learning system may output instructions to a human user via a digital assistant for actions the user should perform at each of multiple time steps. The instructions may be generated, for example, in the form of natural language (communicated as speech and / or on-screen text) based on the actions selected by the reinforcement learning system. The reinforcement learning system selects actions such that the actions contribute to the performance of the task. A monitoring system (e.g., a video camera system) may be provided to monitor for discrepancies (e.g., due to human error) between the actions the reinforcement learning system instructed the user to perform and the actions (if any) the user actually takes at each time step. The reinforcement learning system may use the monitoring system to determine whether the task is completed. During the policy-based training phase and / or other phases in which a history database is generated, experience tuples may record the actions the user actually took based on the instructions, rather than the actions the reinforcement learning system instructed the user to perform. A reward value for each experience tuple may be generated by comparing the actions taken by the user to a corpus of data representing human experts performing the task, for example, using techniques known from imitation learning. Note that if the user performs an action incorrectly (i.e., performs an action different from the action the reinforcement learning system instructed the user to perform), this adds one or more sources of noise to those that may already be present in the environment. During the training process, the reinforcement learning system may identify actions that the user performed incorrectly beyond a certain probability. If an action is identified, the reinforcement learning system may warn the user to be careful when instructing the user to perform such an identified action. Alternatively or additionally, the reinforcement learning system may learn not to instruct the user to perform the identified action, i.e., an action that the user is likely to perform incorrectly.

[0172] More generally, a digital assistant instructing a user may include receiving a request for assistance from the user at the digital assistant and, in response to the request, identifying a series of tasks, e.g., steps or subtasks, of an overall task for the user to perform. The digital assistant can then be used to output instructions for the tasks to be performed, e.g., steps or subtasks, to the user for one or more tasks in the series, e.g., task by task, up to, e.g., the final task in the series. This can be done using natural language, e.g., on the digital assistant's display and / or using a speech synthesis subsystem. Visual observations, e.g., video observations, and / or audio observations of the user performing the task can be captured, e.g., using the digital assistant. Such a system can then be used to determine whether the user successfully completed, e.g., a step or subtask, i.e., a task from the aforementioned answer. If there are more tasks to complete, the digital assistant can responsively proceed to the next task (if any) in the series of tasks, e.g., by outputting instructions for the next task to be performed. In this way, the user can be guided step by step through the series of tasks to perform the overall task. During training of the neural network, training rewards may be generated, for example, from video data representing examples of the overall task (if a corpus of such data is available) or from simulations of the overall task.

[0173] In a further aspect, a digital assistant device including the aforementioned system is provided. The digital assistant may also include a user interface that allows the user to request assistance and output information. In an embodiment, this is a natural language user interface and may include a keyboard, a voice input / output subsystem, and / or a display. The digital assistant may further include an assistance subsystem configured to identify a series of tasks for the user to perform in response to a request. In an embodiment, this may include a generative (large-scale) language model, particularly for dialogue, such as a conversational agent such as Sparrow or Chinchilla. The digital assistant may include an observation capture subsystem for capturing visual and / or audio observations of the user performing the tasks and an interface for the aforementioned language model neural network (which may be implemented locally or remotely). The digital assistant may also include an assistance control subsystem configured to assist the user. The assistance control subsystem may be configured to perform the aforementioned steps, for example, with respect to one or more tasks in the series of tasks, for example, up to the final task in the series of tasks. More specifically, the assistance control subsystem outputs instructions for the task to be performed to the user, uses the observation capture subsystem to capture visual or audio observations of the user performing the task, and determines from such responses whether the user successfully completed the task. In response, the digital assistant may proceed to the next task in the sequence and / or control the digital assistant to, for example, stop capturing observations.

[0174] In some implementations, the environment may not include humans or animals. However, in other implementations, the environment may include humans or animals. For example, the agent may be an autonomous vehicle within an environment (e.g., a geographic location) where humans (e.g., pedestrians or other vehicle drivers / passengers) and / or animals are present, and the autonomous vehicle itself may optionally include humans. The environment may also be at least one room (e.g., a residence) containing one or more people. The humans or animals may be elements of the environment involved in, e.g., modified by, a task (indeed, the environment may consist substantially of humans or animals). For example, the environment may be a medical or veterinary environment including at least one human or animal subject, and the task may relate to performing a medical procedure (e.g., a surgical procedure) on the subject. In further implementations, the environment may include a human user interacting with an agent in the form of an item of user equipment, such as a digital assistant. The item of user equipment provides a user interface between the user and a computer system (the same computer system(s) implementing the reinforcement learning system or a different computer system). The user interface may enable a user to input data to and / or receive data from the computer system, and the agent may be controlled by an action selection policy to perform a communication task relevant to the user, such as providing the user with information about a topic and / or allowing the user to specify components of a task for the computer system to perform. For example, the communication task may be teaching the user a skill such as how to speak a language or how to navigate around a geographic location, or the task may be allowing a user to define a three-dimensional shape to the computer system, such as allowing the computer system to control an additive manufacturing (3D printing) system to fabricate an object having the shape.Actions may include outputting information to the user (e.g., in a particular format, at a particular speed, etc.) and / or configuring an interface to receive input from the user. For example, actions may include setting a problem for the user to tackle related to the skill (e.g., asking the user to select from multiple options regarding correct use of language or asking the user to speak a passage of language aloud) and / or receiving input from the user (e.g., registering the selection of one of the options or using a microphone to record the spoken passage of language). Rewards may be generated based on measurements of how well the task was performed. For example, this may be done by measuring how well the user is learning a topic, e.g., performing an instance of a skill (e.g., as measured by an automated skill assessment unit of the computer system). In this way, a customized learning system may be provided that is tailored to the user's aptitude and current knowledge. In another example, if the communication task is to specify components of a task to be performed by a computer system, the actions may include presenting the user with a user interface (visual, tactile, or auditory) that allows the user to specify the components of the task and receiving user input using the user interface. A reward may be generated based on a measurement of how well and / or easily the user is able to specify the components of the task to be performed by the computer system, e.g., how completely or well a three-dimensional object was specified. This may be determined automatically, or the reward may be specified by the user, e.g., as a subjective measurement of the user experience. In this way, again, a customized system for a user to control a computer system may be provided, tailored to the user's aptitudes and current knowledge.

[0175] Optionally, in any of the above embodiments, the observations at any given time step may include data from previous time steps that may be useful in characterizing the environment, such as actions performed at the previous time step, rewards received at the previous time step, or both.

[0176] 4 is a flow diagram of an exemplary process 400 for training a policy neural network and a value neural network. For convenience, process 400 is described as being performed by a system of one or more computers located at one or more locations. For example, a training system suitably programmed in accordance with this specification, such as training system 100 of FIG. 1, may perform process 400.

[0177] The system may repeatedly perform iterations of process 400 with different imagination paths generated using the world model neural network to update the parameters of the policy neural network and the value neural network.

[0178] The policy neural network is configured to process policy network inputs including updated latent representations according to current values ​​of policy network parameters to generate action selection outputs, and the value neural network is configured to process value network inputs including updated latent representations according to current values ​​of value network parameters to generate predicted state values.

[0179] The system may continue to perform iterations of process 400 until a termination criterion for training the neural network is met, for example, until the parameters converge, until a threshold amount of actual time has elapsed, or until a threshold number of iterations of process 400 have occurred.

[0180] The system receives a latent representation, which exists in a latent space and characterizes the current state of the environment (step 402). The latent representation may be generated by a world model neural network based on a current observation (e.g., an image or video frame) about the current state of the environment. The world model neural network may include a sequence model neural network, an encoder neural network, a decoder neural network, a dynamics predictor neural network, a reward neural network, and a continuation predictor neural network. When the current state is an initial state of the environment, the latent representation may be generated by the encoder neural network based on a predetermined initial hidden state of the sequence model neural network and the current observation.

[0181] The system generates an imaginative path, where the imaginative path includes an updated latent representation for each of a plurality of time steps in the imaginative path (step 404). The imaginative path is initiated from the received latent representation. That is, at the first time step, the updated latent representation is the received latent representation. Briefly, as described in more detail below with reference to FIG. 5, generating each imaginative path may include iteratively generating action selection outputs using a policy neural network, using the action selection outputs to select actions that the agent will perform, and then using a world model neural network to generate updated latent representations based on at least the selected actions.

[0182] For each updated latent representation in the imagination path, the system determines a predicted reward (step 406) and a predicted state value (step 408). The predicted state value represents a predicted value of the environment in the state characterized by the updated latent representation in terms of a cumulative assessment of the predicted reward to be received by the agent. For example, starting from the state characterized by the updated latent representation, the long-term time-adjusted (e.g., time-discounted) total predicted reward that the agent will receive by selecting an action using the action selection output can be determined from the predicted state value.

[0183] More specifically, for each updated latent representation in the imagination pathway, the system may identify a predicted reward by processing the updated latent representation and an updated hidden state of a sequence model neural network according to the current value of a reward network parameter using a reward neural network to generate a predicted reward, which may involve processing the current latent representation and the selected action using a sequence model neural network to generate an updated hidden state. Similarly, the system may identify a predicted state value by processing the updated latent representation using a value neural network to generate a predicted state value.

[0184] For each updated latent representation in the imagination path, the system identifies a corresponding goal state value based on the predicted reward and predicted state value for the updated latent representation in the imagination path (step 410). For example, the goal state value may be identified by calculating a bootstrap γ-return, which integrates both the predicted reward and the predicted state value. An exemplary method for calculating the bootstrap γ-return is described above with reference to FIG. 2B.

[0185] The system identifies updates to the current values ​​of the policy network parameters by optimizing an actor objective function (step 412). In some embodiments, the actor objective function includes a first term that depends on the target state value. In some embodiments, the actor objective function includes a second term that depends on the entropy of the action selection output produced by the policy neural network. Examples of actor objective functions are described above with reference to FIG. 2B.

[0186] The system may determine updates by calculating the gradient of each of the actor objective functions with respect to the policy network parameters, e.g., via backpropagation, and then applying an update rule, e.g., the Adam update rule, the AdamW update rule, the Adagrad update rule, or the RMSProp update rule, to each gradient.

[0187] For each updated latent representation in the imagination path, the system applies a symmetric logarithmic transformation to each goal state value to generate a transformed goal state value (step 414).

[0188] The system encodes each transformed destination state value to generate an encoded transformed destination state value for each updated latent representation in the imagination path (step 416). For example, the system may use a two-hot encoding technique to generate the encoded transformed destination state values, each of which may be represented by a vector of a predetermined length having a predetermined number of non-zero entries.

[0189] The system identifies updates to the current values ​​of the value network parameters by optimizing a critic objective function (step 418). The critic objective function includes terms that assess the difference between the predicted state values ​​and the transcoding target state values. An example of a critic objective function is described above with reference to Figure 2B.

[0190] The system may identify updates by calculating the gradient of each of the critic objective functions with respect to the value network parameters, e.g., via backpropagation, and then applying an update rule, e.g., the Adam update rule, the AdamW update rule, the Adagrad update rule, or the RMSProp update rule, to each gradient.

[0191] 5 is a flow diagram of an exemplary process 500 for generating imaginative pathways of latent representations. For convenience, process 500 is described as being performed by a system of one or more computers located at one or more locations. For example, a training system suitably programmed in accordance with this specification, such as training system 100 of FIG. 1, may perform process 500.

[0192] The system may repeatedly perform process 500 to generate an imagination path that includes an updated latent representation at each of multiple time steps in the imagination path. The imagination path begins with the latent representation received in step 402 of process 400. That is, at the first time step, the updated latent representation is the received latent representation.

[0193] The system processes the current latent representation using the policy neural network according to the current values ​​of the policy network parameters to generate an action selection output (step 502). At the first time step, the current latent representation is the received latent representation. At any subsequent time step, the current latent representation is the updated latent representation generated at the immediately preceding time step.

[0194] When the environment is in the state characterized by the current latent representation, the system selects an action to be performed by the agent using the action selection output (step 504). Selecting an action by using the action selection output is similarly described above with reference to FIG.

[0195] The system generates an updated latent representation based on the current latent representation and the selected action (step 506). If the agent executes the selected action when the environment is in a state characterized by the current latent representation, the state to which the environment may transition is characterized by the updated latent representation.

[0196] To do this, the system processes the current latent representation and the selected action using a sequence model neural network to generate an updated hidden state according to the current values ​​of the sequence model network parameters, and then processes the updated hidden state using a dynamics predictor neural network to generate an updated latent representation.

[0197] Prior to training the policy neural network and the value neural network by executing processes 400 and 500, the system trains the world model neural network to ensure the quality of the imagination pathway. Specifically, the system trains the encoder neural network and the dynamics predictor neural network to promote the generation of high-quality latent representations, such as latent representations that can capture salient information (e.g., task-specific) contained in observations of the environment or that predict subsequent states of the environment. The system also trains the sequence model neural network to more accurately simulate the latent dynamics of the environment, i.e., the state transitions of the environment that will be caused by the agent performing the proposed action; trains the reward neural network to more accurately estimate the reward that the agent will receive from the environment in response to performing the selected action; and trains the continuation predictor neural network to more accurately predict whether the current state is the final state of the environment.

[0198] 6 is a flow diagram of an exemplary process 600 for training a world neural network. For convenience, process 600 is described as being performed by a system of one or more computers located at one or more locations. For example, a training system suitably programmed in accordance with this specification, such as training system 100 of FIG. 1, may perform process 600.

[0199] The system may sample history paths from a replay buffer and train a world model neural network with the sampled history paths. The replay buffer may be any available memory that stores history paths, which are generated as a result of an agent (or other agents) interacting with the environment (or other instances of the environment) to attempt to perform one or more tasks.

[0200] In some implementations, each history path may include a tuple at each of multiple points in time along the history path. At a given point in time, the tuple includes: (i) a current observation characterizing the current state of the environment, (ii) a current action being performed by the agent, (iii) a current reward being received by the agent, and (iv) a ground truth continuation flag indicating whether the current state of the environment is the final state of the environment (e.g., a goal state of the environment achieved by completing a task).

[0201] The system may repeatedly perform iterations of process 600 on different history paths sampled from the replay buffer to update the parameters of the world model neural network. The system may continue to perform iterations of process 600 until a termination criterion for training the world model neural network is met, for example, until the parameters converge, until a threshold amount of wall-clock time has elapsed, or until a threshold number of iterations of process 600 have been performed.

[0202] Typically, the system may perform the following steps 602-612 at each of multiple time steps in the history path.

[0203] The system uses a sequence model neural network to process (i) the current latent representation generated in the immediately preceding time step and characterizing the previous state of the environment, and (ii) data specifying the previous selection action selected in the immediately preceding time step, according to the current values ​​of the sequence model network parameters, to generate an updated hidden state characterizing the current state of the environment (step 602).

[0204] The system processes the updated hidden state and the current observations of the current state of the environment according to the current values ​​of the encoder network parameters using an encoder neural network to generate an updated latent representation that characterizes the current state of the environment (step 604).

[0205] The system processes the updated hidden states using the dynamics predictor neural network according to the current values ​​of the dynamics predictor network parameters to generate estimated updated latent representations (step 606). The estimated updated latent representations characterize the current state of the environment and are estimates of the updated latent representations that may be generated by the encoder neural network.

[0206] The system processes the updated latent representation, and optionally the updated hidden state, using a decoder neural network according to the current values ​​of the decoder network parameters to generate a current observation that reconstructs the current state of the environment (step 608).

[0207] The system processes the updated latent representation, and optionally the updated hidden state, using a reward neural network according to the current values ​​of the reward network parameters to generate a predicted reward (step 610). The predicted reward is a prediction of the reward the agent would receive if the environment were in its current state.

[0208] The system processes the updated latent representation, and optionally the updated hidden state, using a continuation predictor neural network according to the current values ​​of the decoder network parameters to generate a predicted continuation predictor flag indicating whether the current state is the terminal state of the path (step 612).

[0209] The system determines updates to the current values ​​of the parameters of the world model neural network based on optimizing the combined objective function (step 614). Generally, the combined objective function assesses the difference between the output of the neural network and the ground truth information contained in the historical path. An example of a combined objective function is described above with reference to Figure 2A.

[0210] More specifically, determining updates to current values ​​of parameters of the world-model neural network may include determining updates to current values ​​of the sequence model neural network parameters and the dynamics predictor neural network parameters based on calculating gradients of a dynamics objective function with respect to the sequence model neural network parameters and the dynamics predictor neural network parameters. The dynamics objective function relies on the Kullback-Leibler divergence between (i) the latent representations generated by the encoder neural network and (ii) the latent representations generated by the dynamics predictor neural network. Here, the sequence model neural network parameters are parameters of the sequence model neural network configured to process the current latent representations and the selected actions to generate updated hidden states for the sequence model neural network. The dynamics predictor neural network parameters are parameters of the dynamics predictor neural network configured to process the updated hidden states to generate updated latent representations. The encoder neural network is configured to process the observations to generate latent representations that characterize the current state of the environment.

[0211] The system can determine updates by calculating the gradient of each of the combined objective functions with respect to the appropriate network parameters, e.g., through backpropagation, and then applying an update rule, e.g., the Adam update rule, the AdamW update rule, the Adagrad update rule, or the RMSProp update rule, to each gradient.

[0212] During training, the system may incorporate any number of techniques to improve the speed, effectiveness, or both of the training process. For example, the system may refrain from using hyperparameter annealing, weight decay, or dropout. As another example, the system may initialize the values ​​of some parameters of the neural network described above, such as the reward network parameters and value network parameters, to zero. This can accelerate early learning by mitigating problems with delayed start of training caused by large predicted rewards, which may otherwise occur if parameters are randomly initialized at the beginning of training.

[0213] It will be appreciated that the techniques for training an encoder-decoder neural network can be applied to a wider range of tasks, for example, any of which may not involve controlling an agent. As a general example, an encoder-decoder neural network may be configured to receive an input data item, process the input data item to generate a latent representation of the input data item, and process the latent representation of the input data item to generate a reconstruction of the input data item.

[0214] Encoder-decoder neural networks can operate on any of a variety of different types of data items, such as, for example, images, video, or audio signals. As a particular example, encoder-decoder neural networks can operate on high-dimensional (and therefore large) data items. For example, encoder-decoder neural networks can be used to compress and reconstruct high-resolution images, such as medical and geographic images. More specifically, the input data items may include images of real-world entities captured by an imaging sensor, such as a camera or LIDAR sensor, or images of real-world entities captured by a medical imaging sensor, such as an X-ray or MRI machine. The latent representation of the input data items constitutes a compressed version of the input data items that may be stored and / or transmitted, for example, over a communications network. The compressed version of the input data items may be retrieved from storage and / or received over a communications network, and the input data items may be reconstructed.

[0215] 7 is a flow diagram of an exemplary process 700 for training an encoder-decoder neural network. For convenience, process 700 is described as being performed by one or more computer systems located at one or more locations. For example, a training system such as training system 100 of FIG. 1, or other training system including a Simulog transformation engine, can be suitably programmed in accordance with this specification to perform process 700.

[0216] The system receives training data items (step 702).

[0217] The system processes the training data items using an encoder-decoder neural network to determine latent representations of the training data items and generate reconstructions of the training data items according to current values ​​of the parameters of the encoder-decoder neural network (step 704). The latent representations are lower-dimensional, i.e., more compact, versions of the training data items. The encoder-decoder neural network includes an encoder neural network and a decoder neural network.

[0218] More specifically, the system processes training data items using an encoder neural network to determine latent representations of the training data items. The system processes the latent representations of the training data items using a decoder neural network to generate outputs of the encoder-decoder neural network. The system then applies a symmetric exponential transformation to the outputs of the encoder-decoder neural network to generate reconstructions of the training data items. For example, a symmetric exponential transformation may be sign(x)(exp(|x|)-1) or sign(x)(10 |x|-1 ), where x is the output of the encoder-decoder neural network.

[0219] The system applies a symmetric logarithmic transformation to the training data items to generate transformed training data items (step 706). For example, the symmetric logarithmic transformation is of the form sign(x)ln(|x|+1) or sign(x)log(|x|+1), where x is the training data item.

[0220] The system determines updates to the current values ​​of the encoder-decoder network parameters based on evaluating a loss function that assesses the difference between the reconstruction of the training data item and the transformed training data item (step 708). For example, the loss function can be a mean squared error loss function. The system can determine the updates by computing the gradients of each of the combined objective functions with respect to the parameters of the encoder-decoder neural network, e.g., through backpropagation, and then applying an update rule, e.g., the Adam update rule, AdamW update rule, Adagrad update rule, or RMSProp update rule, to each gradient.

[0221] Figure 8 shows quantitative examples of the performance gains that can be achieved by using the neural networks described herein across a number of tasks spanning seven different domains, including continuous and discrete actions, visual and low-dimensional inputs, dense and sparse rewards, different reward scales, 2D and 3D worlds, and procedural generation.

[0222] Tasks include proprioceptive sensor-based control tasks and visual control tasks (described in arXiv:1801.00690), Atari tasks (described in arXiv:1903.00374 and arXiv:1207.4708), behavioral suite tasks, or short b-suite tasks (described in arXiv:1908.03568), and crafter tasks (described in arXiv:2109.06780).

[0223] The vertical axis of Figure 8 shows the return an agent receives when controlled using each of multiple neural network systems, with higher returns indicating better performance. In Figure 8, the DreamerV3 system corresponds to a system including the neural network described herein. The other systems are existing reinforcement learning systems for agent control, such as the MPO system described in arXiv:1806.06920, the DDPG system described in arXiv:1509.02971, and the D4PG system described in arXiv:1804.08617. It will be recognized that DreamerV3 outperforms existing systems in all of these tasks. In particular, the DreamerV3 system is able to achieve this performance gain across these tasks after being trained using the same set of hyperparameters listed above in Table 1.

[0224] This specification uses the term "configured" in the context of systems and computer program components. To say that one or more computer systems are configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that, when running, causes the system to perform the operation or action. To say that one or more computer programs are configured to perform a particular operation or action means that the one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0225] Embodiments of the subject matter and functional operations described herein can be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, such as the structures disclosed herein and their structural equivalents, or one or more combinations thereof. Embodiments of the subject matter described herein may be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium for execution by or controlling the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random-access or serial-access memory device, or one or more combinations thereof. Alternatively or additionally, the program instructions may be encoded in an artificially generated propagated signal, such as a mechanically generated electrical, optical, or electromagnetic signal, generated to encode information for transmission to a receiving device suitable for execution by the data processing apparatus.

[0226] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or computers. An apparatus may also be or include special-purpose logic circuitry, such as an FPGA (field-programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, an apparatus may optionally include code that creates an execution environment for a computer program, such as code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

[0227] A computer program, which may be called or described as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted, declarative or procedural, and can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, for example, one or more scripts stored in a markup language document, in a single file dedicated to the program, or in multiple cooperating files, for example, files that store one or more modules, subprograms, or code portions. A computer program can be deployed to be executed on one computer or to be executed on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communications network.

[0228] As used herein, the term "database" is used broadly to refer to any collection of data. The data need not be structured in any particular way, or even at all, and may be stored on storage devices in one or more locations. Thus, for example, an index database may contain multiple collections of data, each of which may be organized and accessed in a different way.

[0229] Similarly, the term "engine" is used broadly herein to refer to a software-based system, subsystem, or process programmed to perform one or more specific functions. Typically, an engine is implemented as one or more software modules or components installed on one or more computers at one or more locations. In some cases, one or more computers are dedicated to a particular engine, and in other cases, multiple engines may be installed and run on the same computer or computers.

[0230] The processes and logic flows described herein may be implemented by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows may also be implemented by special purpose logic circuitry, such as an FPGA or an ASIC, or a combination of special purpose logic circuitry and one or more programmed computers.

[0231] A computer suitable for running a computer program may be based on a general-purpose or special-purpose microprocessor, or both, or on other types of central processing units. Generally, the central processing unit receives instructions and data from a read-only memory, a random-access memory, or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by, or incorporated in, special-purpose logic circuitry. Generally, a computer also includes one or more mass storage devices, such as magnetic, magneto-optical, or optical disks, for storing data, or is operatively connected to receive data from or transmit data to them, or both. However, a computer need not have such devices. Furthermore, a computer may be incorporated in other devices, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device (e.g., a universal serial bus (USB) flash drive), to name just a few.

[0232] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, including, by way of example, semiconductor memory devices such as EPROM, EEPROM, and flash memory devices, magnetic disks such as internal hard disks or removable disks, magneto-optical disks, and CD-ROM and DVD-ROM disks.

[0233] To provide for interaction with a user, embodiments of the subject matter described herein may be implemented in a computer having a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, as well as a keyboard and pointing device, such as a mouse or trackball, by which the user can provide input to the computer. Other types of devices may also be used to interact with the user. For example, feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback, and input from the user may be received in any form, including acoustic, speech, or tactile input. Furthermore, a computer may interact with a user by sending and receiving documents to a device used by the user, for example, by sending a web page to a web browser on the user's device in response to a request received from the web browser. A computer may also interact with a user by sending text messages or other forms of messages to a personal device (e.g., a smartphone running a messaging application) and then receiving a reply message from the user.

[0234] A data processing device for implementing machine learning models may also include dedicated hardware accelerator units, for example for handling the common and computationally intensive part of machine learning training or production, i.e., inference, workloads.

[0235] The machine learning model can be implemented and deployed using a machine learning framework, for example, the TensorFlow framework or the JAX framework.

[0236] Embodiments of the subject matter described herein may be implemented in a computing system that includes a back-end component, e.g., a data server, or includes a middleware component, e.g., an application server, or includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an embodiment of the subject matter described herein, or includes any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include a local area network (LAN), a wide area network (WAN), e.g., the Internet.

[0237] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server sends data (e.g., HTML pages) to a user device, for example, to display the data to and receive user input from a user interacting with the device functioning as a client. Data generated at the user device, e.g., results of user interaction, can be received from the device by the server.

[0238] While the specification contains many specific implementation details, these should not be construed as limiting the scope of any invention or the scope of patentable subject matter, but rather as descriptions of features that may be unique to particular embodiments of a particular invention. Certain features described herein as separate embodiments may also be implemented in combination in a single embodiment. Conversely, various features described in a single embodiment may also be implemented in multiple embodiments, either separately or in any suitable subcombination. Furthermore, even if features may be described above as functioning in a particular combination and originally claimed as such, one or more features from a claimed combination may, in some cases, be deleted from the combination, and the claimed combination may be directed to subcombinations or variations of the subcombinations.

[0239] Similarly, while operations are illustrated in the figures and described in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order or sequential order shown, or that all of the operations shown be performed, to achieve desirable results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the program components and systems described may generally be integrated into a single software product or packaged into multiple software products.

[0240] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims may be performed in a different order and still produce desirable results. As an example, the processes depicted in the accompanying figures do not necessarily require the particular order shown or sequential order to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. 1. A method for training a policy neural network having a plurality of policy network parameters used to select an action to be performed by an agent in an environment, comprising: receiving a latent representation in a latent space that characterizes a current state of the environment; At each of several time steps, processing the current latent representation using the policy neural network and according to the current values ​​of the policy network parameters to generate an action selection output; using the action selection output to select an action that will be performed by the agent when the environment is in a state characterized by the current latent representation; and processing the current latent representation and the selected action to generate an updated latent representation that characterizes a state to which the environment would transition if the agent performed the selected action when the environment was in the state characterized by the current latent representation; generating an imagination path of a latent representation starting with the received latent representation by performing For each latent representation in the imagination pathway, Determining the expected reward; and processing the latent representation using a value neural network having a plurality of value network parameters and according to current values ​​of the value network parameters to generate a predicted state value representing a time-adjusted sum of predicted rewards that would be received by the agent by selecting an action using an action selection output starting from the state characterized by the latent representation; determining a goal state value for each latent representation in the imagination pathway from the predicted reward, and the predicted state value for the latent representation in the imagination pathway; determining updates to the current values ​​of the policy network parameters by optimizing an actor objective function that includes a first term that depends on the target state value; applying a symmetric logarithmic transformation to each goal state value to generate a transformed goal state value for each latent representation in the imagination pathway; encoding each transformed destination state value to generate, for each latent representation in the imagination pathway, an encoded transformed destination state value; determining updates to the current values ​​of the value network parameters by optimizing a critic objective function that assesses differences between the predicted state values ​​and the transcoding target state values; A method comprising:

2. 2. The method of claim 1, wherein the symmetric logarithmic transformation is of the form sign(x)ln(|x|+1) or sign(x)log(|x|+1), where x is the target state value.

3. 3. The method according to claim 1, wherein the transcoding destination state value is represented by a vector of a predetermined length having a predetermined number of non-zero entries.

4. generating the predicted state value for each latent representation in the imagination pathway, processing the latent representations to assign a respective probability to each predetermined interval in a sequence of predetermined intervals between possible ranges of state values; generating an output of the value neural network by calculating a weighted average of the state values ​​represented by the sequence of predetermined intervals weighted by the respective probabilities; and 4. The method of claim 1, further comprising applying a symmetric exponential transformation to the output of the value neural network to generate the predicted state value for the latent representation in the imagination pathway.

5. The symmetric exponential transformation is sign(x)(exp(|x|)-1) or sign(x)(10 ｜ｘ｜－１ 5. The method of claim 4, wherein x is the output of the value neural network.

6. Determining the updates to the current values ​​of the policy network parameters by optimizing the actor objective function comprises: generating a downscaled goal state value for each latent representation in the imagination pathway from the goal state value by using a scaling factor; and 6. The method of claim 1, comprising determining the update to the current value of the policy network parameter by optimizing the actor objective function that includes the first term that depends on the downscaled target state value.

7. generating the downscaled target state value 7. The method of claim 6, comprising dividing the goal state value by a maximum given integer value and the scaling factor, the scaling factor being calculated using a percentile of each of one or more goal state values ​​for the latent representation in the imagination pathway.

8. 8. The method of claim 1, wherein the actor objective function includes a second term that depends on the scaled entropy of the action selection output produced by the policy neural network, the scaled entropy being scaled by a predetermined entropy scaling factor.

9. processing the current latent representation and the selected action to generate the updated latent representation includes: processing the current latent representation and the selected action using a sequence model neural network according to current values ​​of sequence model network parameters to update a hidden state of the sequence model neural network to generate an updated hidden state that characterizes the current state of the environment; and 9. The method of claim 1, comprising processing the updated hidden state using a dynamics predictor neural network according to current values ​​of dynamics predictor network parameters to generate the updated latent representation.

10. receiving the latent representation characterizing the current state of the environment; 10. The method of claim 9, comprising processing observations characterizing the current state of the environment using an encoder neural network according to current values ​​of encoder network parameters to generate the latent representation.

11. processing the latent representation and the updated hidden state using a decoder neural network according to current values ​​of decoder network parameters to generate a reconstruction of the observation; applying the symmetric logarithmic transformation to the observations to generate transformed observations; and 11. The method of claim 10, further comprising: determining updates to the current values ​​of the decoder network parameters based on calculating a gradient of a prediction objective function that includes a first term that assesses a difference between the reconstruction of the observation and the transformed observation, and with respect to the decoder network parameters.

12. Determining the predicted reward for each latent representation in the imagination pathway includes:

12. The method of claim 1, comprising processing the latent representation and updated hidden states of the sequence model neural network using a reward neural network according to current values ​​of reward network parameters to generate the predicted reward.

13. applying the symmetric logarithmic transformation to rewards received by the agent when the environment is in the state characterized by the latent representation to generate transformed rewards; and 13. The method of claim 12, further comprising determining an update to the current value of the reward network parameter based on calculating a gradient of the predicted objective function including a second term that assesses a difference between the predicted reward and the transformed reward and with respect to the reward network parameter.

14. 14. The method of claim 1, further comprising determining updates to the current values ​​of the sequence model neural network parameters and the dynamics predictor neural network parameters based on calculating gradients of a dynamics objective function with respect to the sequence model neural network parameters and the dynamics predictor neural network parameters, the dynamics objective function depending on a Kullback-Leibler divergence between (i) the latent representations generated by the encoder neural network and (ii) the latent representations generated by the dynamics predictor neural network.

15. the encoder neural network and the decoder neural network are each configured as a respective convolutional neural network or a multi-layer perceptron (MLP); the dynamics predictor neural network and the reward neural network are each configured as a multi-layer perceptron (MLP); 15. The method of claim 1, wherein the sequence model neural network is configured as a recurrent neural network.

16. The convolutional neural network, the multilayer perceptron (MLP), or the recurrent neural network is The method of claim 15 , comprising one or more normalization layers and one or more sigmoidal linear unit (SiLU) activation layers.

17. 17. The method of claim 15, further comprising initializing the current values ​​of the reward network parameters and the value network parameters to zero.

18. 18. The method of claim 1, further comprising using the policy neural network to select an action to be performed by the agent in the environment by processing inputs derived from data characterizing the current state of the environment at the time step, using the policy neural network to generate an action selection output used to determine an action to be performed by the agent at the time step.

19. 1. A method executed by one or more computers for training an encoder-decoder neural network, the encoder-decoder neural network being configured to receive an input data item, process the input data item to generate a latent representation of the input data item, and process the latent representation of the input data item to generate a reconstruction of the input data item; The method comprises: receiving a training data item; processing the training data items using the encoder-decoder neural network to determine latent representations of the training data items and to generate reconstructions of the training data items according to current values ​​of encoder-decoder network parameters; applying a symmetric logarithmic transformation to the training data items to generate transformed training data items; and determining updates to the current values ​​of the encoder-decoder network parameters based on evaluating a loss function that assesses a difference between the reconstruction of the training data item and the transformed training data item.

20. 19. The method of claim 18, wherein the symmetric logarithmic transformation is of the form sign(x) ln(|x|+1) or sign(x) log(|x|+1), where x is the training data item.

21. 20. The method of any one of claims 18 to 19, wherein the loss function is a mean squared error loss function.

22. generating the reconstruction of the input data items, processing the training data items to determine the latent representations of the training data items; processing the latent representations of the training data items to generate outputs of the encoder-decoder neural network; and 21. A method according to any one of claims 18 to 20, comprising applying a symmetric exponential transformation to the output of the encoder-decoder neural network to generate the reconstruction of the training data items.

23. The symmetric exponential transformation is sign(x)(exp(|x|)-1) or sign(x)(10 ｜ｘ｜－１ 22. The method of claim 21, wherein x is the output of the encoder-decoder neural network.

24. 23. A method according to any one of claims 18 to 22, wherein the input data items comprise images.

25. 25. A system comprising one or more computers and one or more storage devices storing instructions, the instructions, when executed by the one or more computers, operable to cause the one or more computers to perform the respective method operations of any one of claims 1 to 24.

26. 26. A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the respective method operations of any one of claims 1 to 25.

Images