Training reinforcement learning agents to perform multiple tasks across diverse fields

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training a policy neural network for selecting actions to be performed by an agent interacting with an environment. In one aspect, a method includes receiving a potential representation characterizing a current state of an environment; generating a potentially represented imaginary trajectory; for each potential representation in the imaginary trajectory: determining a predicted reward; and generating a predicted state value; determining a target state value for each potential representation; determining an update to the current value of the policy network parameter; applying a symmetric logarithmic transformation to each target state value; encoding each transformed target state value to generate an encoded transformed target state value; and determining an update to the current value of the value network parameter by optimizing the reviewer objective function.

Description

Cross Reference to Related Applications

[0001] 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 and is incorporated by reference into the disclosure of this application. BACKGROUND

[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 a current state of the environment.

[0004] Some reinforcement learning systems select an action to be performed by an agent in response to receiving a given observation in accordance with an output of a neural network.

[0005] A neural network is a machine learning model that employs one or more layers of nonlinear units to predict an output for a received input. Some neural networks are deep neural networks that include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as 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 a received input in accordance with current values of a respective set of parameters. SUMMARY

[0006] This specification describes a system implemented as computer programs on one or more computers in one or more locations that trains a policy neural network that can be used to control an agent interacting with an environment by selecting actions to be performed by the agent. The system implements a general-purpose reinforcement learning framework that can efficiently train the policy neural network to select actions that enable the agent to accomplish a large number of tasks across a broad range of domains.

[0007] These tasks can include, for example, navigating to a specified location in an environment, identifying a particular object in an environment, manipulating a particular object in a specified manner, and so on. In some cases, the tasks are long-horizon tasks that require the agent to select and perform thousands or millions of actions in sequence in order to accomplish the task. In some cases, the tasks are tasks that are performed by the agent in a complex environment, e.g., a complex workspace that includes a large number of obstacles and / or target objects, a large number of target locations, or both.

[0008] The domain of each task generally defines properties and / or characteristics of various aspects of the task, e.g., the environment in which the agent is deployed, the space of possible actions that the agent can perform, the data (also referred to in this specification as “observations”) that is received by the agent that characterizes the environment, etc. The domains of these tasks can differ from one another in one or more of these aspects, e.g., in terms of the action space (e.g., continuous vs. discrete action space), the observation space (e.g., high-dimensional visual input vs. low-dimensional textual input, or 2D vs. 3D camera images), the data budget, the reward frequency, or the reward scale.

[0009] In one aspect, this specification describes a reinforcement learning method performed by one or more computers for training a policy neural network to control an agent acting in an environment based on rewards received as a result of performed actions. The policy neural network can be trained using a reinforcement learning objective function that depends on the received rewards, and the rewards or values depending on the rewards can be scaled or “squashed” using a symmetric log transformation as described herein.

[0010] In general, one innovative aspect of the subject matter described in this specification can be embodied in a method of training a policy neural network having a plurality of policy network parameters for selecting actions for an agent to perform in an environment, the method including: receiving a latent representation, the latent representation being in a latent space and characterizing a current state of the environment; generating an imagined trajectory of latent representations, starting from the received latent representation, by performing the following at each of a plurality of time steps: processing a current latent representation using the policy neural network to generate an action selection output in accordance with current values of the policy network parameters; using the action selection output to select an action that the agent would perform 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, the updated latent representation characterizing 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; for each latent representation in the imagined trajectory: determining a predicted reward; and processing the latent representation using a value neural network having a plurality of value network parameters to generate a predicted state value in accordance with current values of the value network parameters, the predicted state value representing a time-adjusted sum of predicted rewards that the agent would receive by using action selection outputs to select actions starting from a state characterized by the latent representation; determining a target state value for each latent representation in the imagined trajectory in accordance with the predicted reward and the predicted state values for the latent representations in the imagined trajectory; determining an update 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 values; applying a symmetric log transform to each target state value to generate a transformed target state value for each latent representation in the imagined trajectory; encoding each transformed target state value to generate an encoded transformed target state value for each latent representation in the imagined trajectory; and determining an update to the current values of the value network parameters by optimizing a critic objective function that measures a difference between the predicted state values and the encoded transformed target state values.

[0011] The symmetric log transform can have the following form: or where x is the target state value.

[0012] The encoded transformed target state value can be represented by a vector of a predetermined length having a predetermined number of non-zero entries.

[0013] Generating the predicted state value for each latent representation in the imagined trajectory can include processing the latent representation to assign a respective probability for each predetermined interval in a predetermined sequence of intervals between a range of possible state values, generating an output of the value neural network by computing a weighted average of state values represented by the predetermined sequence of intervals weighted by the respective probabilities, and 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 imagined trajectory.

[0014] The symmetric exponential transformation can have the form: or where x is the output of the value neural network.

[0015] Determining the update to the current values of the policy network parameters by optimizing the actor objective function can include generating, for each latent representation in the imagined trajectory, a scaled-down target state value from the target state value using a scaling factor, and determining the update to the current values of the policy network parameters by optimizing the actor objective function including a first term dependent on the scaled-down target state value.

[0016] Generating the scaled-down target state value can include dividing the target state value by a maximum of a given integer value and the scaling factor, where the scaling factor is computed using respective percentiles of one or more target state values of the latent representations in the imagined trajectory.

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

[0018] Processing the current latent representation and the selected action to generate the updated latent representation can include 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 representing the current state of the environment, and processing the updated hidden state using a dynamic predictor neural network according to current values of dynamic predictor network parameters to generate the updated latent representation.

[0019] Receiving the latent representation representing the current state of the environment can include processing an observation representing 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 can further include processing the latent representation and the updated hidden state using a decoder neural network in accordance with current values of decoder network parameters to generate a reconstruction of the observation; applying the symmetric log transform to the observation to generate a transformed observation; and determining, with respect to the decoder network parameters, updates to the current values of the decoder network parameters based on computing gradients of a predictive objective function, the predictive objective function including a first term that measures a difference between the reconstruction of the observation and the transformed observation.

[0021] Determining, for each latent representation in the imagined trajectory, the predicted reward can include processing the latent representation and the updated hidden state of the sequence model neural network using a reward neural network in accordance with current values of reward network parameters to generate the predicted reward.

[0022] The method can further include applying the symmetric log transform to a reward received by the agent when the environment is in a state characterized by the latent representation to generate a transformed reward; and determining, with respect to the reward network parameters, updates to the current values of the reward network parameters based on computing gradients of the predictive objective function, the predictive objective function including a second term that measures a difference between the predicted reward and the transformed reward.

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

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

[0025] The convolutional neural network, the multi-layer perceptron (MLP), or the recurrent neural network can include one or more layer normalization layers and one or more sigmoid linear unit (SiLU) activation layers.

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

[0027] The method can further include using the policy neural network to select an action for the agent to perform in the environment by processing, using the policy neural network, an input derived from data characterizing a current state of the environment at a time step to generate an action selection output for determining an action for the agent to perform at the time step.

[0028] In another aspect, the specification describes a method performed by one or more computers for training an encoder-decoder neural network. The encoder-decoder neural network can 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 a latent representation of the input data item (including the encoder network output of the latent representation), and process the latent representation of the input data item (using the decoder neural network) to generate a reconstruction of the input data item (at the decoder network output). An update to current values of the encoder-decoder network parameters is determined based on evaluating a loss function that measures a difference between the reconstruction of the training data item and a transformed training data item that has been transformed using a symmetric log transform.

[0029] In another aspect, the specification describes a method performed by one or more computers for training an encoder-decoder neural network. The encoder-decoder neural network can 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 a latent representation of the input data item (including the encoder network output of the latent representation), and process the latent representation of the input data item (using the decoder neural network) to generate a reconstruction of the input data item (at the decoder network output). An update to current values of the encoder-decoder network parameters is determined based on evaluating a loss function that measures a difference between the reconstruction of the training data item and a transformed training data item that has been transformed using a symmetric log transform.

[0030] The symmetric log transform can have the form: or where x is the training data item.

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

[0032] Generating the reconstruction of the input data item can 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 the encoder-decoder neural network, and applying a symmetric exponential transform to the output of the encoder-decoder neural network to generate the reconstruction of the training data item.

[0033] The symmetric exponential transform can have the form: or where x is the output of the encoder-decoder neural network.

[0034] The input data item can include an image.

[0035] Other embodiments of this aspect include corresponding computer systems, apparatus, and computer programs recorded on one or more computer storage devices, each configured to perform the actions of the methods. A system of one or more computers can be so configured by virtue of software, firmware, hardware, or any combination thereof installed on the system that in operation cause the system to perform the actions. One or more computer programs can be so configured by virtue of including instructions that, when executed by data processing apparatus, cause the apparatus to perform the actions.

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

[0037] Many complex tasks, such as robotic tasks, involve reward scales that can vary across domains of application, between particular tasks, and across stages of learning, sometimes by orders of magnitude. This can impair stability and can result in slow or failed learning. The described techniques can implement a general-purpose reinforcement learning (RL) framework that trains a policy neural network and one or more additional neural networks based on a latent representation of an environment to achieve robust performance across multiple tasks, even including long-horizon tasks (requiring an agent to select and execute thousands or millions of actions in sequence) and tasks performed in sparse-reward or difficult-to-explore environments, across a wide range of domains.

[0038] The described reinforcement learning framework can use a fixed set of hyperparameters that can be easily adapted to these diverse domains and thus are easy to set when a new task needs to be learned. Under the described framework, training an instance of the policy neural network for each new task can thus consume less clock time and fewer computational resources, e.g., memory and processing power, than conventional RL systems because hyperparameter tuning by iteratively running multiple training trials is no longer needed. The described reinforcement learning framework is also scalable, meaning that it is capable of training arbitrarily large policy neural networks where increased model size, e.g., in terms of model parameters, corresponds to higher task performance and improved data efficiency.

[0039] In particular, by utilizing a symmetric logarithmic transformation of the magnitudes of large positive and negative output values of the compressed neural network while preserving the input sign, the described framework mitigates common problems that can arise when using some of the existing loss functions and normalization techniques, including divergence, stagnation, and non-stationarity, and instead ensures stable training of the neural network across diverse domains by unifying gradient signals of different orders of magnitude.

[0040] The details of one or more embodiments of the subject matter described in this specification are set forth in the accompanying drawings and the following description. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. BRIEF DESCRIPTION OF DRAWINGS

[0041] Figure 1 An example training system is shown.

[0042] Figure 2A is an example illustration of training a world model neural network.

[0043] Figure 2B is an example illustration of training a policy neural network and a value neural network.

[0044] Figure 3 An example data flow for selecting an action to be performed by an agent when interacting with an environment is shown.

[0045] Figure 4 is a flow diagram of an example process for training a policy neural network and a value neural network.

[0046] Figure 5 is a flow diagram of an example process for generating imagined trajectories.

[0047] Figure 6 is a flow diagram of an example process for training a world model neural network.

[0048] Figure 7is a flowchart of an example process for training an encoder-decoder neural network.

[0049] Figure 8 Quantitative examples of performance gains achievable by using the neural networks described in this specification are shown.

[0050] Like reference numbers and designations in the various drawings indicate like elements. DETAILED DESCRIPTION

[0051] This specification describes a training system implemented as a computer program on one or more computers in one or more locations for training a policy neural network. During or after training, the policy neural network can be used to control an interaction of a reinforcement learning agent with an environment by processing, at each time step in a plurality of time steps, an input (i.e., an “observation”) derived from data that characterizes a current state of the environment at the time step to generate an output for determining an action to be performed by the agent.

[0052] In particular, the training system implements a general-purpose reinforcement learning (RL) framework to train the policy neural network by generating a plurality of imagined trajectories and then training the policy neural network on the plurality of imagined trajectories. By doing so, the training system trains the policy neural network such that the policy neural network can be used to control the agent to perform any of a plurality of tasks.

[0053] Each imagined trajectory includes data that characterizes a predicted sequence of state transitions of the environment that would result from the agent performing actions while performing a task. To this end, each imagined trajectory includes a plurality of latent representations in a latent space that characterize different states of the environment across a plurality of time steps during the imagined trajectory. The plurality of latent representations includes a latent representation that characterizes an initial state of the environment, a latent representation that characterizes a terminal state of the environment (e.g., a goal state of the environment achieved after completion of the task), and one or more latent representations that respectively characterize intermediate states between the initial state and the terminal state of the environment.

[0054] In this specification, 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 includes an ordered set of numerical values, e.g., a tensor (e.g., a vector or a matrix) of numerical values, that typically has a predetermined dimensionality that is lower than a dimensionality of an observation. In this specification, the space of possible tensors having the predetermined dimensionality is referred to as a “latent space.”

[0055] Figure 1An example training system 100 is shown. The training system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

[0056] The training system 100 includes a policy neural network 120, a value neural network 130, and a world model neural network 140. The policy neural network 120 is a neural network configured to receive, at a given time step during a trajectory, a policy network input that includes a latent representation characterizing a current state of the environment at the given time step ; and process the policy network input according to a set of parameters (referred to in this specification as “policy network parameters”) to generate an action selection output that can be used to select an action to be performed by the agent at the given time step . The latent representation is also referred to later as an updated latent representation; obtaining the updated latent representation at the given time step is described further below. The action selection output and how it can be used to select an action are described next in several examples. In one example, the action selection output can include a respective numerical probability value for each action in a set of possible actions that can be performed by the agent. If used to select an action, the action to be performed by the agent can be selected, for example, by sampling the actions according to their probability values, or by selecting the action with the highest probability value. In another example, the action selection output can directly define the action to be performed by the agent, for example, by defining values that should be applied to torques of the joints of the robotic agent. Such an action can be an action that is predicted to result in the highest long-term, time-discounted reward if performed by the agent in response to the observation.

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

[0058]

[0059]

[0060] ​​​​In this example, the Q-value of an action is an estimate of the “reward” that would result from the agent performing that action in response to the current observation, and thereafter selecting future actions performed by the agent based on the current values ​​of the policy network parameters.

[0061] A reward is a cumulative measure of the "rewards" received by an agent, e.g., the time-discounted sum of rewards. The agent receives a reward at each time step, where the reward is specified by a scalar value and represents, for example, the agent's progress toward completing an assigned task.

[0062] The value neural network 130 is the following neural network : The neural network is configured to, at a given time step during the trajectory The value network input is received at the given time step. The updated potential representation of the current state at ; and according to a set of parameters (referred to in this specification as "value network parameters") process the value network inputs to generate predicted state values, i.e., predicted values ​​of the current state of the environment. In some implementations, the predicted state value represents the agent's potential representation after updating from the environment. The time-adjusted sum of predicted rewards that will be received during the trajectory starting from the current state by selecting actions using the action selection output of the policy neural network.

[0063] The world model neural network 140 captures the training system 100's understanding of the environment ("world"), allowing the world model neural network 140 to make predictions about trajectories in latent space based on its understanding of the past. In general, the world model neural network 140 is configured to simulate the environmental state transitions that would result from actions performed by the agent that might be selected at multiple points during the trajectory using the action selection outputs of the policy neural network 120.

[0064] World model neural network 140 may include sequence model neural network 150, encoder-decoder neural network (which in turn includes 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.

[0065] The sequence model neural network 150 is a neural network configured to receive a sequence model network input at a given time step during a trajectory, the sequence model network input comprising a current latent representation characterizing a previous state of the environment and data specifying a previously selected action in accordance with a set of parameters (referred to in this specification as “sequence model network parameters”) to update a hidden state of the sequence model neural network , for example, by processing the previous latent representation and the previously selected action to modify a current state that has been generated by processing another previous latent representation and data specifying another previously selected action . The subscript is generally used herein to refer to parameters of the world model neural network 140, and of any neural network that the world model neural network comprises. A “parameter” as used herein refers to a trainable parameter of a neural network, such as a weight, as distinct from a hyperparameter.

[0066] At a given time step , the hidden state of the sequence model neural network 150 after processing the world model network input will be referred to in this specification as an updated hidden state , the updated hidden state characterizing a current state of the environment. The current state is the state that the environment would transition to if the agent performed the previously selected action when the environment was in the previous state characterized by the current latent representation . For an initial state of the environment, the updated hidden state of the sequence model neural network 150 can be a predetermined initial hidden state.

[0067] Like a latent representation, a “hidden state” refers to a representation of an environment state that is more data-efficient, e.g., more compact, than an observation of the same environment state. A hidden state comprises an ordered set of numerical values—e.g., a tensor (e.g., a vector or a matrix) of numerical values—typically having a predetermined dimensionality that is lower than that of the observation.

[0068] The encoder neural network 160 comprised in the encoder-decoder neural network is a neural network configured to receive an encoder network input at a given time step at a given time step, receive an encoder network input comprising an updated hidden state characterizing a current state of the environment and a current observation of the current state of the environment ; and process the encoder network input to generate an updated latent representation characterizing the current state of the environment in accordance with a set of parameters (referred to in this specification as “encoder network parameters”) .

[0069] While observations of an environment state typically comprise high-dimensional sensor data characterizing the environment state, e.g., image or video data, a latent representation characterizing the same environment state can be represented as an ordered set of numerical values, e.g., a numerical vector or matrix, typically having a lower dimensionality than the observation. For example, each observation can be an RGB image having 64 x 64 x 3 pixels, while each latent representation can be a 32-dimensional array comprising 32 latent variables.

[0070] The decoder neural network 165 included in the encoder-decoder neural network is a neural network configured to, at a given time step , receive a decoder network input comprising an updated latent representation characterizing a current state of the environment and, in some implementations, an updated hidden state characterizing the current state of the environment ; and process the decoder network input to generate a reconstructed current observation of the current state of the environment in accordance with a set of parameters (referred to in this specification as “decoder network parameters”) . The reconstructed current observation is a reconstruction of the current observation received by the encoder neural network 160 .

[0071] The dynamic predictor neural network 170 is a neural network configured to, at a given time step , receive a dynamic predictor network input comprising an updated hidden state characterizing a current state of the environment ; and process the dynamic predictor network input to generate an estimated updated latent representation characterizing the current state of the environment in accordance with a set of parameters (referred to in this specification as “dynamic predictor network parameters”) . The estimated updated latent representation is an estimate of the updated latent representation characterizing the current state of the environment and generatable by the encoder neural network 160 .

[0072] Estimated updated latent representation can have the same potential representation as the updated The same dimensions. Unlike the encoder neural network, the dynamics predictor neural network 170 is configured to be independent of the current observation of the current state of the environment. Generate an updated latent representation from limited information in the case of , i.e., in the implementation, the dynamic predictor network input does not include the observation .

[0073] The reward neural network 180 is the following neural network : The neural network is configured to At , receiving a reward network input, the reward network input includes an updated latent representation representing the current state of the environment , and in some implementations includes an updated hidden state representing the current state of the environment ; and according to a set of parameters (referred to as “reward network parameters” in this specification) processes the reward network input to generate the predicted reward . Reward for prediction is the reward the agent will receive when the environment is in its current state As used in this paper, A neural network may refer to any neural network that makes predictions about rewards, dynamics, continuations, or observations, for example (see below).

[0074] The continuation predictor neural network 190 is a neural network as follows : The neural network is configured to At , receiving a continuation predictor network input, the continuation predictor network input including an updated latent representation representing the current state of the environment , and in some implementations, includes an updated hidden state representing the current state of the environment ; and according to a set of parameters (referred to in this specification as "continuation predictor network parameters") to process the continuation predictor network input to generate the predicted trajectory continuation flag , the predicted trajectory continuation flag indicates whether the current state is a terminal state in the trajectory, that is, whether another updated latent representation should be generated to represent the future state of the environment.

[0075] In some implementations, the predicted trajectory continuation flag includes binary values, where 1 defines that the current state is a terminal state in the trajectory (such that no further updated latent representation should be generated), and 0 defines that the current state is not a terminal state in the trajectory (such that another updated latent representation should be generated), or vice versa. For example, a trajectory can terminate when the agent completes a task, or after a threshold number of time steps, etc.

[0076] The policy neural network 120, the value neural network 130, the sequence model neural network 150, the encoder neural network 160, the decoder neural network 165, the dynamics predictor neural network 170, the reward neural network 180, and the continuation predictor neural network 190 can each be implemented with any appropriate neural network architecture that enables them to perform their described functions.

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

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

[0079] In some implementations, the value neural network 130 is configured to generate the predicted state value with a discrete regression approach. Specifically, the value neural network 130 first generates a probability distribution over a sequence of predetermined intervals using, for example, a softmax layer of the value neural network, where may be any positive integer (e.g., 63, 127, 255, etc.). That is, the probability distribution is a sequence of predetermined intervals where each predetermined interval of the sequence of predetermined intervals The corresponding probability scores are assigned. Then, the value neural network 130 generates a predicted state value by calculating the average value of the state values ​​represented by the sequence of predetermined intervals. The average value is calculated as a weighted average weighted by the corresponding probability scores.

[0080] In some implementations, the 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. When there is more than one recurrent layer, the hidden state may be, for example, the hidden state of the last recurrent layer in a stack of multiple recurrent layers included in the sequence model neural network 150.

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

[0082] In some implementations, similar to the encoder neural network 160 and the decoder neural network 165, the dynamics predictor neural network 170, the reward neural network 180, and the continuation predictor neural network 190 are each configured as a respective multilayer perceptron (MLP) including one or more fully connected layers.

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

[0084] Depending on their architecture, the above outputs can be generated in any of a number of ways. For example, the updated hidden state can be generated deterministically by, for example, the last recurrent layer of a sequence model neural network: .

[0085] As another example, the updated latent representation , the estimated updated latent representation , reconstructed current observation , predicted rewards or track continuation sign Each of one or more of can be randomly generated. For example, in some implementations, the output of the encoder neural network can be used to represent the updated latent representation Parameterize the distribution to sample from: ; The output of the decoder neural network can be used to reconstruct the current observation Parameterize the distribution to sample from: The output of the dynamic predictor neural network will be the updated latent representation from which the estimate Parameterize the distribution to sample from: ; The output of the reward neural network will be used to predict the reward Parameterize the distribution to sample from: ; and the output of the continuation predictor neural network will be used to predict the trajectory continuation flag Parameterize the distribution to sample from: .

[0086] The training system 100 includes or accesses a replay buffer 104 that stores historical trajectories that may have been generated as a result of an agent (or another agent) interacting with an environment (or another instance of an environment) to 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.

[0087] In some implementations, each history track may include a history track period to A tuple at each of the multiple time points between. At a given time point The tuple consists of: (i) the current observation representing the current state of the environment , (ii) the current action performed by the agent , (iii) the current reward received by the agent , and (iv) a true-valued continuation flag indicating whether the current state of the environment is the final state of the environment (e.g., the intended state of the environment achieved upon completion of the task) .

[0088] The training engine 110 of the training system 100 repeatedly selects a historical trajectory from the replay buffer 140 and trains the world model neural network 140 based on the selected historical trajectory using a suitable world model learning technique as described below. After training the world model neural network 140, the training engine 110 trains the policy neural network 120 and the value neural network 130 using the (trained) world model neural network 140. This is achieved by generating a plurality of imagined trajectories using the (trained) world model neural network 140 and then jointly training the policy neural network 120 and the value neural network 130 based on the plurality of imagined trajectories using a suitable reinforcement learning technique.

[0089] Training the policy neural network 120 and the value neural network 130 based on the plurality of imagined trajectories means that the policy neural network 120 and the value neural network 130 are each configured to receive a latent representation as input during training (e.g., rather than directly processing actual observations characterizing an environment state, e.g., sensor data such as image or video data).

[0090] Implementations of the training system 100 include applying a symlog transformation to the input A symlog transformation engine 112 applies a symlog transformation to the input to generate a transformed output, and is symmetric about the origin (zero value of the input and the log-transformed input) while preserving the sign of the input or variable The log can have any base. As one example, the symlog transformation can take the following form: where x is the input or variable. As one example, the symlog transformation can use a log with base 10 (or with another number as the base), and can for example take the following form:

[0091] The symlog transformation compresses the magnitude of large positive and negative values while preserving the sign of the input. By utilizing such a symlog transformation, the training engine 110 mitigates common problems that can arise when using some of the existing loss functions and normalization techniques, including divergence, stagnation, and non-stationarity, and instead ensures stable training of the neural networks across diverse domains by unifying gradient signals of different orders of magnitude.

[0092] Figure 2A is an example illustration of training a world model neural network. For example, the world model neural network can be trained using a symlog transformation engine 112 that applies a symlog transformation to the input Figure 1corresponding world model neural network 140, which can include a sequence model neural network 150, an encoder-decoder neural network (which in turn includes an encoder neural network 160 and a decoder neural network 165), a dynamics predictor neural network 170, a reward neural network 180, and a continuation predictor neural network 190.

[0093] The training engine 110 performs the operations shown in FIG. 2 for each of the plurality of historical trajectories selected from the replay buffer 104 to train the neural networks using suitable world model learning techniques to enable the neural networks to more accurately model the effects of actions performed by the agent when interacting with the environment. Figure 2A

[0094] As shown, at each of a plurality of time steps during a given historical trajectory selected from the replay buffer 104, the sequence model neural network receives a sequence model network input that includes a current latent representation characterizing a previous state of the environment and data specifying a previously selected action; and processes the sequence model network input to generate an updated hidden state characterizing a current state of the environment. The current state of the environment is the state that the environment would transition to if the agent performed the previously selected action when the environment was in the previous state characterized by the current latent representation.

[0095] For example, at time step , the sequence model neural network generates an updated hidden state characterizing a current state of the environment. The current state of the environment at time step is the state that the environment would transition to if the agent performed the previously selected action when the environment was in the previous state characterized by the current latent representation .

[0096] At each of a plurality of time steps during a given historical trajectory selected from the replay buffer 104, an encoder neural network (“enc”) receives an encoder network input that includes an updated hidden state characterizing a current state of the environment and a current observation of the current state of the environment; and processes the encoder network input to generate an updated latent representation characterizing the current state of the environment.

[0097] For example, at time step ​​At this point, the encoder neural network (“enc”) generates an updated latent representation that represents the current state of the environment. The environment at time step The current state is when the environment is in the state represented by the current potential The previous state of the representation is if the agent performed the previously selected action The state that the environment will transition to.

[0098] At multiple time steps during a given historical trajectory selected from the replay buffer 104 At each time step in , the decoder neural network (“dnc”) receives the decoder network input, which includes the updated latent representation representing the current state of the environment. and in some implementations includes an updated hidden state representing the current state of the environment ; and processes the decoder network input to produce a reconstructed current observation that represents the current state of the environment . Refactoring of current observations is the current observation received by the encoder neural network 160 Reconstruction.

[0099] For example, at time step At , the decoder neural network (“dec”) generates a reconstruction of the current observation The current observation of this reconstruction is the current observation received by the encoder neural network (“enc”) Reconstruction.

[0100] At multiple time steps during a given historical trajectory selected from the replay buffer 104 At each time step in , the dynamics predictor neural network receives the dynamics predictor network input, which includes the updated hidden state representing the current state of the environment. ; and processes the dynamics predictor network input to generate an updated latent representation that represents an estimate of the current state of the environment The estimated updated latent representation is an updated latent representation that represents the current state of the environment and can be generated by the encoder neural network 160 estimated value.

[0101] For example, at time step At , the dynamics predictor neural network generates an updated latent representation that represents an estimate of the current state of the environment The estimated updated latent representation is the updated latent representation The updated potential representation characterizes the environment at time step The current state at and can be generated by the encoder neural network 160.

[0102] At multiple time steps during a given historical trajectory selected from the replay buffer 104 At each time step in , the reward neural network receives a reward network input that includes an updated latent representation representing the current state of the environment. , and in some implementations includes an updated hidden state representing the current state of the environment ; and processes the decoder network input to generate predicted rewards . Reward for prediction is the reward the agent will receive when the environment is in its current state For example, at time step At , the reward neural network 180 generates a predicted reward , the predicted reward is the reward for the environment at time step The reward the agent will receive when in the current state predictions.

[0103] At multiple time steps during a given historical trajectory selected from the replay buffer 104 At each time step in , the continuation predictor neural network receives a continuation predictor network input consisting of an updated latent representation representing the current state of the environment. , and in some implementations includes an updated hidden state representing the current state of the environment ; and processes the continuation predictor network input to generate the predicted trajectory continuation flag ,The predicted trajectory continuation flag indicates whether the current state is a terminal state in the trajectory.

[0104] For example, at time step At , the continuation predictor neural network generates the predicted trajectory continuation flag , the predicted trajectory continuation flag indicates that the current state is not the terminal state in the historical trajectory. Therefore, for example, the training system continues to use the sequence model neural network to generate an updated hidden state representing the future state of the environment The future state of the environment is when the environment is in the state represented by the current potential If the agent performs the selected action when representing the current state The environment will transition to the state.

[0105] For a given historical trajectory selected from the replay buffer 104, the training engine 110 updates 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 based on optimizing a prediction objective function .

[0106] The prediction objective function may include a first term that measures, for each time step of a plurality of time steps during the given trajectory , a difference (e.g., a squared difference) between (i) a reconstructed current observation and (ii) a transformed current observation generated by a symmetric log transform engine 112 included in the training engine 110 from a current observation included in the historical trajectory .

[0107] In particular, to generate the transformed current observation, the symmetric log transform engine 112 applies a symmetric log transform to the current observation . For example, the symmetric log transform can take the following form: or where is the current observation. The symmetric log transform can be applied to any type of observation, including, for example, low-dimensional inputs. When applied to a vector or a matrix, it can be applied to each element of the vector or matrix.

[0108] The first term trains the decoder neural network to generate reconstructed current observations that are close to the current observations included in the historical trajectory, and, by means of backpropagation, trains the encoder neural network to generate semantically meaningful updated latent representations and trains the sequence model neural network to generate semantically meaningful updated hidden states.

[0109] The prediction objective function may include a second term that measures, for each time step of a plurality of time steps during the given trajectory , a difference (e.g., a squared difference) between a predicted reward and a transformed reward generated by the symmetric log transform engine 112 included in the training engine 110 from a reward included in the historical trajectory . The second term trains the reward neural network to generate more accurate predicted rewards, i.e., predicted rewards that are close to the rewards included in the historical trajectory, and, by means of backpropagation, trains the encoder neural network to generate semantically meaningful updated latent representations and trains the sequence model neural network to generate semantically meaningful updated hidden states.​

[0110] the prediction objective function may comprise a third term that measures, for each time step of a plurality of time steps during a given trajectory, the difference between the predicted trajectory continuation flag and the true value trajectory continuation flag. The third term trains the continuation predictor neural network to generate more accurate predicted trajectory continuations, i.e. predicted trajectory continuation flags that are close to the trajectory continuation flags included in the history trajectory, and, by means of backpropagation, trains the encoder neural network to generate semantically meaningful updated latent representations and trains the sequence model neural network to generate semantically meaningful updated hidden states.

[0111] For example, the prediction objective function may be expressed as: where, in fact, the terms are each computable as a difference in cross-entropy loss (binary classification loss for the continuation term).

[0112] For a given trajectory selected from the replay buffer 104, the training engine 110 further updates the values of the parameters of the dynamic predictor neural network included in the world model neural network 140 based on optimizing the prediction objective function .

[0113] The dynamic objective function may comprise a term that depends on a difference that measures, for each time step of a plurality of time steps during a given trajectory, (i) the updated latent representation generated by the encoder neural network and (ii) the estimated updated latent representation generated by the dynamic predictor neural network . This term trains the dynamic predictor neural network to generate an estimated updated latent representation that is close to the updated latent representation generated by the encoder neural network.

[0114] For example, the difference can be determined as a Kullback-Leibler (KL) divergence, and the dynamic objective function may be: where sgd denotes a stop-gradient operator applied to the parameters of the encoder neural network (i.e. during training, gradients do not backpropagate into the encoder neural network), and where ​Refers to taking the maximum of the variable inside the brackets. Taking the maximum effectively eliminates the influence of the loss once it drops below a minimum level (1 nat), allowing training to focus on other losses, such as the prediction loss.

[0115] For a given trajectory selected from the replay buffer 104, the training engine 110 performs a training operation based on the representation of the target function Optimize to update the parameters of the encoder neural network included in the world model neural network value.

[0116] With dynamic objective function Same, represents the objective function May include terms that depend on the difference between the two for multiple time steps during a given trajectory. At each time step in , we measure (i) the updated latent representation generated by the encoder neural network and (ii) the estimated updated latent representation generated by the dynamic predictor neural network This term trains the encoder neural network to generate an updated latent representation that is close to the estimated updated latent representation generated by the dynamic predictor neural network.

[0117] For example, the difference can be determined as the Kullback-Leibler (KL) divergence and represents the objective function It can be: where sg represents the stopped gradient operator applied to the parameters of the dynamic predictor neural network.

[0118] 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 is trained by optimizing a combined objective function that combines the above-mentioned prediction objective function , dynamic objective function and represents the objective function For example, the combined objective function can be: .

[0119] In this example, 、 and is applied to the prediction objective function , dynamic objective function and represents the objective function In the implementation, 、 and are hyperparameters for training system 100 (some illustrative hyperparameter values ​​are given later).

[0120] Figure 2B is an example diagram of training policy neural network and value neural network. For example, policy neural network and value neural network can be respectively Figure 1 The policy neural network 120 and the value neural network 130 correspond to each other. 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 imaginary trajectories generated using the world model neural network 140.

[0121] The training engine 110 performs a training on each of the plurality of imagined trajectories using a suitable reinforcement learning technique using the world model neural network and the policy neural network and the value neural network. Figure 2B The operations shown enable the policy neural network to generate action selection outputs that can be used to select actions that cause the agent to effectively perform one or more tasks.

[0122] In some implementations, each imagined trajectory may include an imagined trajectory period. to A tuple at each of the multiple time points between. At a given time point The tuple includes: (i) an updated latent representation representing the current state of the environment , (ii) the updated hidden state representing the current state of the environment , (iii) the current action selected by the action selection output (iii) Predicted current reward , and (iv) a continuation flag indicating whether the current state of the environment is the predicted terminal state of the environment (e.g., the intended state of the environment achieved after completing the task) .

[0123] As shown in the figure, the imagined trajectory includes to The updated latent representation at each time step in the multiple time steps between For the first time step , the updated latent representation The encoder neural network is based on the predetermined initial hidden state of the sequence model neural network and based on the current observation included in the historical trajectory representing the initial state of the environment For any subsequent time step, the updated latent representation is generated by the dynamic predictor neural network based on the updated hidden state of the sequence model neural network. The updated latent representation of the generated estimate .

[0124] Thus, imagined trajectories differ from historical trajectories stored in replay buffer 104 in that each imagined trajectory includes at least one updated latent representation that is not generated based on observations of the state of the environment, e.g., sensor data such as image or video data. Instead, the imagined trajectory includes an updated latent representation generated based on the hidden states of the sequence model neural network.

[0125] In summary, the 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.

[0126] At multiple time steps during a given imagined trajectory At each time step in , the policy neural network receives the policy network input, which includes the representation of the environment at time step The updated potential representation of the current state at ; and processes the policy network input to generate an action selection output , the action selection output can be used to select the agent at time step Action to be performed ,in , and can represent imaginary model states. For example, at time step At this point, the policy neural network generates an action selection output , the action selection output can be used to select the action at time step Select an action .

[0127] At multiple time steps during a given imagined trajectory At each time step in , the value neural network receives the value network input, which includes the representation of the environment at time step The updated potential representation of the current state at ; and processes the value network input to generate a predicted state value ( ). For example, at time step At time step , the value neural network generates an environment The predicted value of the current state (exist Figure 2B It is called " ”).

[0128] In which the value neural network generates a value between the range of possible state values a predetermined interval a sequence of a probability distribution In implementations in which the output is implemented as a value, the training engine 110 can use the symmetric logarithmic transform engine 112 to apply a symmetric exponential transform to the output of the value neural network to generate the predicted state value. Generally, the symmetric exponential transform is the inverse operation of the symmetric logarithmic transform described previously.

[0129] For example, the symmetric exponential transform may take the form: or where x is the output of the value neural network. The predicted state value at each time step in a given imagined trajectory can be computed as:

[0130] For a given imagined trajectory, the training engine 110 determines a target state value at each time step and updates the values of the parameters of the policy neural network based on optimizing an actor target function. The actor target function includes a first term that depends on the target state value. The actor target function trains the policy neural network to generate action selection outputs that can facilitate selecting more effective actions, which in turn causes the value neural network to compute higher predicted state values and causes the reward neural network to compute higher predicted rewards.

[0131] The target state value at each time step may be determined based on the predicted reward generated by the reward neural network, the predicted state value generated by using the value neural network, or both the predicted reward and the predicted state value ( ). For example, the target state value can be determined by computing a bootstrap return that integrates the predicted reward and the predicted state value: , , where is a discount factor, is a trace decay parameter, for which a larger value gives more credit to rewards that are more temporally distant, and is the imagined model state after the trajectory of T time steps. In implementations, and is a hyperparameter of the training system 100. In other approaches, n-step returns can be used.

[0132] In some implementations, a first term included in the actor target function depends on the scaled target state values to better accommodate sparse rewards. Specifically, the training engine 110 scales down larger target state values without scaling up smaller target state values. This is achieved by dividing each target state value by the maximum of (i) a given integer value (e.g., one, two, ten, etc.) and (ii) a scaling factor .

[0133] For example, the scaling factor may be computed using the corresponding quantile, e.g., the percentile of one or more target state values of the latent representation in the trajectory. As a particular example, the scaling factor may be computed as an exponentially decaying average of the range from its 5th to its 95th percentile batch: . Such approaches can accelerate exploration under sparse rewards without sacrificing final performance under dense rewards. A fixed entropy scale is used at the same time.

[0134] In some implementations, the actor target function includes a second entropy regularization term that depends on the entropy of the action selection output generated by the policy neural network. The entropy may be a scaled entropy scaled by a predetermined entropy scaling factor .

[0135] For example, the actor target function can be: where sgdenotes a stop-gradient operator applied to the parameters of the world model neural network. Here, the max operator only implements scaling when the return exceeds a minimum threshold (e.g., 1), which promotes the use of a single entropy scale (defined by the parameters ) across dense rewards and sparse rewards.

[0136] For a given imagined trajectory, the training engine 110 updates the values of the parameters of the value neural network based on optimizing the critic target function. The critic target function includes a term that measures the difference between (i) the predicted state value and (ii) the encoded transformed target state value generated by the symmetric log transform engine 112 from the target state values. The critic target 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 sum of predicted rewards that the agent will receive during the given imagined trajectory.

[0137] In particular, to generate the encoded transformed target state value, the training engine 110 applies a symmetric logarithmic transformation to the target state value for each time step in the imagined trajectory using the symmetric logarithmic transformation engine 112 to generate a transformed target state value for that time step in the imagined trajectory, and encodes the transformed target state value using a suitable encoding technique to generate an encoded transformed target state value for that time step in the imagined trajectory.

[0138] As explained above, the symmetric logarithmic transformation can for example take the form: or where x is the target state value.

[0139] In implementations, the encoded transformed target state value can be represented by a vector of a predetermined length having a predetermined number of non-zero entries. For example, the training engine can use a duohot encoding technique. The duohot encoding technique allows the value neural network to maintain a fixed distribution over potential returns to avoid wide return distributions that would otherwise slow down training. In this example, the encoded transformed target state value is represented by a vector of length where all elements are zero except for the two entries at positions and The sum of the two entries is 1, and the entry closer to the encoded numerical value is given a greater weight:

[0140] For example, the critic objective function can be: , , where sg denotes a stop gradient operator applied to the parameters of the world model neural network, and the predicted state value the difference between the encoded transformed target state value is computed as a (categorical) cross-entropy loss.

[0141] Referring back to Figure 1 To assist in training the policy neural network 120, the value neural network 130, and the world model neural network 140, the training system 100 maintains predetermined values of a set of hyperparameters 118, for example in one or more physical or logical storage devices accessible to the training system. Hyperparameters are parameters of a neural network that are set before the neural network is trained. Hyperparameters can have an impact on the quality of the training of the neural network, and so setting optimal values of the hyperparameters can be an important part of the training.

[0142] In some implementations, these predetermined values can be adjusted prior to the start of training, e.g., through an automated search process or manual hyperparameter tuning. However, in other implementations, these predetermined values can be fixed, and the training system 100 repeatedly applies this set of hyperparameters with the same predetermined values across different tasks on which the neural networks are to be trained. In those other implementations, the predetermined values of the set of hyperparameters 118 maintained by the training system 100 reduce the need for hyperparameter tuning, which can be time-consuming and computationally costly, and enable it to more efficiently train neural networks across a range of different task domains.

[0143] An example of this set of hyperparameters and their predetermined values is shown in Table 1 below. In this table, “number of latent values” refers to the dimensionality of each latent representation (it is a 32-dimensional array that specifies a probability distribution for each of a set of 32 possible latent variables), “time horizon” refers to the number of time steps included in each imagined trajectory, and EMA is an abbreviation for exponential moving average.

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

[0145] Once trained, some or all of the neural networks included in the training system 100 can be deployed and used to control an agent to perform one or more tasks by interacting with an environment. For example, the sequence model encoder neural network 150 and the encoder neural network 160 can operate in coordination to generate an updated latent representation for each new observation. For example, the sequence model encoder neural network 150 can update the hidden state of the sequence model neural network to generate an updated hidden state that characterizes the current state of the environment. The encoder neural network 160 can process an observation that characterizes the current state of the environment to generate a latent representation . The policy neural network 120 can then process the latent representation of the observation to generate an action selection output that is used to select an action to be performed by the agent in response to the new observation.

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

[0147] Figure 3 An example data flow for selecting actions 308 for the agent 302 to perform when interacting with the environment 304 at each of a plurality of time steps using the policy neural network 320, the sequence model neural network 350, and the encoder neural network 360 is shown.

[0148] The policy neural network 320, the sequence model neural network 350, and the encoder neural network 360 can respectively correspond to the policy neural network 120, the sequence model encoder neural network 150, and the encoder neural network 160 of Figure 1 The trained values of the parameters 114 can be determined (and output) by the training system 100 of Figure 1 .

[0149] At each time step, the sequence model neural network 350 processes a sequence model network input that includes the current observation 306 according to the trained values of the parameters of this sequence model neural network to generate an updated hidden state 352 that characterizes the current state of the environment 304.

[0150] The encoder neural network 360 processes an encoder network input that includes the updated hidden state 352 and the current observation 306 according to the trained values of the parameters of this encoder neural network to generate an updated latent representation 362 that characterizes the current state of the environment 304.

[0151] The policy neural network 320 processes a policy network input that includes the updated latent representation 362 according to the trained values of the parameters of this policy neural network to generate an action selection output 322 that can be used to select an action 308 for the agent 302 to perform in response to the current observation 306. Examples of action selection outputs and how they can be used to select actions have been described above with reference to Figure 1 .

[0152] In particular, during inference and as shown in the example of Figure 3 , each received observation has a one-to-one correspondence with an updated hidden state or an updated latent representation. Note that this is different from training with latent imagination, in which each observation typically has a one-to-many correspondence with updated hidden states or updated latent representations, i.e., during training, a sequence of multiple updated hidden states or updated latent representations is generated from a single observation.

[0153] In some implementations, the environment is a real-world environment, the agent is a mechanical (or electromechanical) agent that interacts with the real-world environment, such as a robot or an autonomous or semi-autonomous land, air, or sea vehicle that operates in or navigates through the environment, and the actions are actions taken by the mechanical agent in the real-world environment to perform a task. For example, the agent can be a robot that interacts with the environment to accomplish a particular task, such as locating or manipulating an object of interest in the environment or moving an object of interest to a specified location in the environment or navigating to a specified destination in the environment.

[0154] In these implementations, the observations can include, for example, one or more of: images, object position data, and sensor data to capture observations as the agent interacts with the environment, such as sensor data from images, distance or position sensors, or from actuators. For example, in the case of a robot, the observations can 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., gravity-compensated torque feedback), and global or relative poses of articles held by the robot. In the case of a robot or other mechanical agent or vehicle, the observations can similarly include one or more of: positions, linear or angular velocities, forces, torques, or accelerations of one or more parts of the agent, and global or relative poses. The observations can be defined in 1, 2, or 3 dimensions, and can be absolute and / or relative observations. The observations can also include, for example, sensed electrical signals such as motor currents or temperature signals; and / or image or video data captured by, for example, cameras or LIDAR (light detection and ranging) sensors, e.g., data from sensors of the agent or data from sensors positioned separately from the agent in the environment.

[0155] In these implementations, the actions can be control signals for controlling a robot or other mechanical agent, e.g., torques for the joints of a robot or higher-level control commands; or can be control signals for controlling an autonomous or semi-autonomous land, air, sea vehicle, e.g., torques for controlling surfaces or other control elements, e.g., steering control elements of the vehicle, or higher-level control commands. The control signals can include, e.g., position, velocity, or force / torque / acceleration data for one or more joints of a robot or portions of another mechanical agent. The control signals can also or instead include electronic control data, such as motor control data, or more generally, data for controlling one or more electronic devices within the environment, the control of which has an effect on the state of the observed environment. For example, in the case of an autonomous or semi-autonomous land or air or sea vehicle, the control signals can define actions for controlling navigation, e.g., steering, and movement, e.g., braking and / or acceleration of the vehicle.

[0156] In some implementations, the environment is a simulation of a real-world environment as described above, and the agents are implemented as one or more computers that interact with the simulated environment. For example, the simulated environment can be a simulation of a robot or vehicle, and the reinforcement learning system can be trained on the simulation, and then, once trained, used in the real world.

[0157] In some implementations, the environment is a real-world manufacturing environment for manufacturing a product, such as a chemical, biological, or mechanical product or a food product. As used herein, “manufacturing” a product also includes refining a starting material to create a product, or processing a starting material, e.g., to remove contaminants, to generate a cleaned or recycled product. The manufacturing plant can 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 are movable between the manufacturing units during manufacturing of the product, e.g., via pipes or mechanical conveyance. As used herein, manufacturing of a product also includes manufacturing of a food product by a kitchen robot.

[0158] The agents can include electronic agents configured to control the manufacturing units or machines, such as robots, that operate to manufacture the product. That is, the agents can include control systems configured to control the manufacturing of a chemical, biological, or mechanical product. For example, the control systems can be configured to control one or more of the manufacturing units or machines, or to control movement of intermediate versions or components of the product between the manufacturing units or machines.

[0159] As one example, the tasks performed by the agent may include tasks of manufacturing a product or an intermediate version or component thereof. As another example, the tasks performed by the agent may include tasks of controlling (e.g., minimizing) the use of resources, such as controlling the consumption of electricity or water or any materials or consumables used in the manufacturing process.

[0160] Actions may include control actions to control the use of a machine or manufacturing cell for processing solid or liquid materials to manufacture a product or an intermediate product or component thereof, or to control the movement of an intermediate version or component of a product within a manufacturing environment, such as between manufacturing cells or machines. In general, an action may be any action that has an effect on the state of an observed environment, such as an action configured to adjust any of the sensed parameters described below. These actions may include actions to adjust the physical or chemical conditions of a manufacturing cell, or actions to control the movement of mechanical parts of a machine or joints of a robot. Actions may include actions that impose operating conditions on a manufacturing cell or machine, or actions that cause a change in a setting to adjust, control, or turn on or off the operation of a manufacturing cell or machine.

[0161] In some implementations, the reinforcement learning scenario includes a reward calculation unit for generating a reward (e.g., in the form of a number), which is typically generated based on an observation. The reward can be associated with a measure of the performance of the task. For example, in the case of a task of manufacturing a product, the metric can include a measure of the number of products manufactured, the quality of the product, the speed at which the product is produced, or a measure of the physical cost of performing the manufacturing task, such as a measure of the amount of energy, material, or other resources used to perform the task. In the case of a task of controlling resource usage, the metric can include any measure of the use of the resource. In the case of a task of controlling a mechatronic agent such as a robot to perform object manipulation, the reward can indicate whether the object has been correctly manipulated according to predefined criteria.

[0162] Generally, the observations of the state of the environment can include any electronic signals representative of the operation of electronic and / or mechanical equipment. For example, the representations 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 manufacturing units or machines, or sensors that sense the movement of materials between manufacturing units or machines. As some examples, such sensors can be configured to sense mechanical movement or force, pressure, temperature; electrical conditions such as current, voltage, frequency, impedance; quantities, levels, flow / movement rates, or flow / movement paths of one or more materials; physical or chemical conditions, e.g., physical state, shape, or configuration or chemical state such as pH; configuration of units or machines such as mechanical configuration of units or machines or valve configuration; image or video sensors for capturing image or video observations of manufacturing units or machines or movement; or any other suitable type of sensor. In the case where the agent is a machine such as a robot, the observations from the sensors can include observations of 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 of an object held or handled by the machine or robot. The observations can also include, e.g., sensed electronic signals such as motor current or temperature signals, or image or video data from cameras or LIDAR sensors, e.g., mounted on the machine. Such sensors can be part of the agent in the environment or located separately from the agent.

[0163] In some implementations, the environment is a real-world environment of a service facility including multiple pieces of electronic equipment, such as a server farm or data center, e.g., a telecommunications data center, or a computer data center for storing or processing data, or any service facility. The service facility can also include auxiliary control equipment, e.g., environmental control equipment, such as temperature control, e.g., cooling, equipment, or air flow control or air conditioning equipment, that controls the operating environment of the equipment. The tasks can include tasks to control, e.g., minimize, the use of resources, such as tasks to control power consumption or water consumption. The agent can include an electronic agent configured to control the operation of the equipment or the operation of the auxiliary, e.g., environmental, control equipment.

[0164] Generally, the actions can be any actions that have an impact on the observed state of the environment, e.g., actions configured to adjust any of the sensed parameters described below. These actions can include actions that control or impose operating conditions on the equipment or auxiliary control equipment, e.g., actions that cause a change of settings to adjust, control, or turn on or off the operation of the equipment or auxiliary control equipment.

[0165] Generally, the observations of the state of the environment can include any electronic signals representative of the operation of the facility or equipment in the facility. For example, the representations of the state of the environment can be derived from observations by any sensors that sense the state of the physical environment of the facility or from observations by any sensors that sense the state of one or more of the equipment or one or more ancillary control equipment. These include sensors configured to sense electrical conditions such as current, voltage, power, or energy; temperature of the facility; fluid flow, temperature, or pressure within the facility or within a cooling system of the facility; or physical facility configuration such as whether a vent is open.

[0166] The reward can be related to a performance metric of the task related to efficient operation of the facility. For example, in the case of a task that controls (e.g. minimizes) usage of a resource, such as in the case of a task that controls usage of electricity or water, the metric can include any measure of usage of the resource.

[0167] 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 power farm or a wind power farm). The task can include a control task that controls power generated by the facility (e.g. controls delivery of power to a power distribution grid, e.g. to meet demand or reduce risk of mismatch between elements of the grid, or to maximize power generated by the facility). The agent can include an electronic agent configured to control generation of power by the facility or coupling of generated power into the grid. The actions can include actions that control electrical or mechanical configuration of a power generator (such as electrical or mechanical configuration of one or more renewable power generation elements) (e.g. control configuration of a wind turbine or solar panel or mirror, or electrical or mechanical configuration of a rotating power generation machine). The mechanical control actions can for example include actions that control a transformation of energy input to electrical energy output (e.g. transformation efficiency or degree of coupling of energy input to electrical energy output). The electrical control actions can for example include actions that control one or more of voltage, current, frequency, or phase of generated power.

[0168] The reward can be related to a performance metric of the task related to power distribution. For example, in the case of a task that controls delivery of power to a power distribution grid, the metric can relate to a measure of power transmitted, or to a measure of electrical mismatch (such as voltage, current, frequency, or phase mismatch) between the power generation facility and the grid, or to a measure of power or energy loss in the power generation facility. In the case of a task that maximizes delivery of power to a power distribution grid, the metric can relate to a measure of power or energy transmitted to the grid, or to a measure of power or energy loss in the power generation facility.

[0169] Generally, the observations of the state of the environment can include any electronic signals representative of the electrical or mechanical operation of the power generation equipment in the power generation facility. For example, the representation of the state of the environment can be derived from observations by any sensors that sense the physical or electrical state of equipment in the power generation facility that is generating power, or the physical environment of such equipment, or the condition of ancillary equipment that supports the power generation equipment. Thus, such observations can include observations of wind levels or solar irradiance, or of local time, date, or season. Such sensors can 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 physical configuration of the equipment; and observations of electrical conditions of the power grid, for example from local or remote sensors. The observations of the state of the environment can also include one or more predictions about 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 power grid.

[0170] As another example, the environment can be a chemical synthesis or protein folding environment, such that each state is a respective state of a protein chain or one or more intermediate or precursor chemicals, and the agent is a computer system for determining how to fold the protein chain or synthesize the chemical. In this example, the actions are possible folding actions for folding the protein chain or actions for assembling the precursor chemicals / intermediates, and the results to be achieved can include, for example, folding the protein such that the protein is stable, and such that it achieves a particular biological function, or providing an efficient synthesis route for the chemical. As another example, the agent can be a mechanical agent that indirectly performs or controls the protein folding actions or chemical synthesis steps, for example by controlling synthesis steps that are automatically selected by the system without human intervention. The observations can include direct or indirect observations of the state of the protein or chemical / intermediate / precursor, and / or can be derived from a simulation. Thus, the system can be used to automatically synthesize a protein with a particular function, such as a protein with a binding site shape that binds a ligand with sufficient affinity to obtain its biological effect that it can be used as a drug. For example, it can be an agonist or antagonist of a receptor or enzyme; or it can be an antibody configured to bind to an antibody target, such as a viral coat protein, or a protein expressed on a cancer cell, for example to act as an agonist of a particular receptor or to prevent binding of another ligand and thus activation of a related biological pathway.

[0171] In a similar manner, the environment can be a drug design environment, such that each state is a respective state of a potential pharmaceutical active compound, and the agent is a computer system for determining elements of the pharmaceutical active compound and / or a synthetic pathway for the pharmaceutical active compound. For example, in a simulation, a drug / synthetic can be designed based on a reward derived from a target for a pharmaceutical active compound. As another example, the agent can be a mechanical agent that performs or controls synthesis of a pharmaceutical active compound.

[0172] In some further applications, the environment is a real-world environment, and the agent manages allocation of tasks across computing resources, e.g., on a mobile device and / or in a data center. In these implementations, the actions can include assigning a task to a particular computing resource. In these applications, the observations can include observations of computing resources such as computing and / or memory capacity or Internet accessible resources; and the actions can include assigning a task to a particular computing resource. The reward can be configured to maximize or minimize one or more of: utilization of computing resources, power, bandwidth, and speed of computation.

[0173] As a further example, the actions can include presenting an advertisement, the observations can include ad impressions or click-through counts or rates, and the reward can characterize prior selections of items or content by one or more users.

[0174] In some cases, the observations can include textual or spoken instructions provided to the agent by a third party (e.g., an operator of the agent). For example, the agent can be an autonomous vehicle, and a user of the autonomous vehicle can provide textual or spoken instructions to the agent (e.g., to navigate to a particular location).

[0175] As another example, the environment can be an electrical, mechanical, or mechatronic design environment, e.g., in which a design of an electrical, mechanical, or mechatronic entity is simulated. The simulated environment can be a simulation of a real-world environment in which the entity is intended to work. The task can be to design the entity. The observations can include observations characterizing the entity, i.e., observations of the mechanical shape or electrical, mechanical, or mechatronic configuration of the entity, or observations of parameters or properties of the entity. The actions can include actions that modify the entity, e.g., actions that modify one or more of the observations. The rewards can include one or more measures of performance of the design of the entity. For example, the rewards can relate to one or more physical properties of the entity, such as weight or strength; or to one or more electrical properties of the entity, such as a measure of efficiency of performing a particular function for which the entity is designed. The design process can include outputting the design for manufacturing, e.g., in the form of computer executable instructions for manufacturing the entity. The process can include fabricating the entity according to the design. Thus, the design of the entity can be optimized, e.g., by reinforcement learning, and then the optimized design is outputted for manufacturing the entity, e.g., as computer executable instructions; and then the entity with the optimized design can be manufactured.

[0176] As previously mentioned, the environment can be a simulated environment. Generally, in the case of a simulated environment, the observations can include simulated versions of one or more of the previously described observations or observation types, and the actions can include simulated versions of one or more of the previously described actions or action types. For example, the simulated environment can be a motion simulation environment, e.g., a driving simulation or a flight simulation, and the agent can be a simulated vehicle that navigates through the motion simulation. In these implementations, the actions can be control inputs that control the simulated user or the simulated vehicle. Generally, the agent can be implemented as one or more computers that interact with the simulated environment.

[0177] The simulated environment can be a simulation of a particular real-world environment and agent. For example, the system can be used to select actions in the simulated environment during training or evaluation of the system, and after training or evaluation or both are complete, the action selection policy can be deployed for controlling a real-world agent in the particular real-world environment that was the subject of the simulation. This can avoid unnecessary wear and tear on the real-world environment or real-world agent and damage to it, and can allow the control neural network to be trained and evaluated in situations that rarely occur or are difficult or unsafe to recreate in the real-world environment. For example, the system can be partially trained using a simulation of a mechanical agent in a simulation of a particular real-world environment, and then deployed to control a real mechanical agent in the particular real-world environment. Thus, in such a case, the observations of the simulated environment are related to the real-world environment, and the selected actions in the simulated environment are related to actions to be performed by the mechanical agent in the real-world environment.

[0178] In some implementations, as described above, the agent can not include a human (e.g., it is a robot). Rather, 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 a task can then be obtained from the digital assistant, and the digital assistant can be used to instruct the user based on the task.

[0179] For example, the reinforcement learning system can output, via a digital assistant, to a human user instructions for actions for the user to perform at each of a plurality of time steps. The instructions can be generated in natural language (sent as sound and / or text on a screen), for example, based on actions selected by the reinforcement learning system. The reinforcement learning system selects these actions so that they help perform the task. A monitoring system (e.g., a camera system) can be provided to monitor the actions actually performed by the user at each time step (if any) in case the action is different from the action the reinforcement learning system instructed the user to perform (e.g., due to human error). Using the monitoring system, the reinforcement learning system can determine whether the task has been completed. During the policy training phase and / or another phase of generating the history database, the experience tuples can record the actions actually performed by the user based on the instructions, rather than the actions the reinforcement learning system instructed the user to perform. The reward value for each experience tuple can be generated, for example, by comparing the actions taken by the user to a corpus of data showing human experts performing the task, for example, using techniques known from imitation learning. Note that if the user does not perform the action correctly (i.e., performs an action different from the action the reinforcement learning system instructed the user to perform), this adds one more source of noise to the sources of noise already present in the environment. During the training process, the reinforcement learning system can identify actions that the user performs incorrectly with a probability that exceeds a certain probability. If so, the reinforcement learning system can warn the user to be careful when the reinforcement learning system instructs the user to perform such identified actions. Alternatively or additionally, the reinforcement learning system can learn not to instruct the user to perform the identified actions, i.e., the actions that the user has a tendency to perform incorrectly.

[0180] More generally, the digital assistant instructing the user can comprise receiving an assistance request from the user at the digital assistant, and determining a series of tasks, e.g. steps of an overall task or subtasks, that the user is to perform in response to the request. Then, for one or more of the tasks in the series, e.g. for each task, e.g. up to a final task in the series, the digital assistant can be used to output an indication to the user of the task, e.g. step or subtask, to be performed. This can be done using natural language, e.g. on a display and / or using a speech synthesis subsystem of the digital assistant. Visual, e.g. video, and / or audio observations of the user performing the task can be captured, e.g. using the digital assistant. The system as described above can then be used to determine whether the user has successfully achieved the task, e.g. step or subtask, i.e. in terms of the answer as described earlier. If there are further tasks to be completed, the digital assistant can respond by proceeding to the next task in the series, if any, e.g. by outputting an indication of the next task to be performed. In this way, the user can be guided step-by-step through a series of tasks to perform an overall task. During training of the neural network, the training reward can be generated, e.g. from video data representing examples of the overall task, if such a corpus of data is available, or from a simulation of the overall task.

[0181] In a further aspect, there is provided a digital assistant apparatus comprising a system as described above. The digital assistant can further comprise a user interface to enable a user to request assistance and output information. In implementations, this is a natural language user interface, and can comprise a keyboard, a speech input output subsystem and / or a display. The digital assistant can further comprise an assistance subsystem configured to determine a series of tasks that the user is to perform in response to the request. In implementations, this can comprise a generative (large) language model, in particular for dialogue, e.g. a conversational agent such as Sparrow or Chinchilla. The digital assistant can have an observation capture subsystem for capturing visual and / or audio observations of the user performing the task; and an interface for the language model neural network as described above (which can be implemented locally or remotely). The digital assistant can also have an assistance control subsystem configured to assist the user. The assistance control subsystem can be configured to perform the steps described above for one or more of the tasks in the series, e.g. up to a final task in the series. More specifically, the assistance control subsystem outputs an indication to the user of the task to be performed, captures visual or audio observations of the user performing the task using the observation capture subsystem, determines from the answer described above whether the user has successfully achieved the task. In response, the digital assistant can proceed to the next task in the series and / or control the digital assistant, e.g. to stop capturing observations.

[0182] In some implementations, the environment can not include humans or animals. However, in other implementations, the environment can include humans or animals. For example, the agent can be an autonomous vehicle, the environment of which is a location (e.g., geographic location) in which humans (e.g., pedestrians or drivers / passengers of other vehicles) and / or animals are present, and the autonomous vehicle itself can optionally house humans. The environment can also be at least one room (e.g., at a residence) that houses one or more humans. The humans or animals can be elements of the environment that are involved in the task - e.g., modified by the task - (indeed, the environment can consist essentially of humans or animals). For example, the environment can be a medical or veterinary environment that includes at least one human or animal subject, and the task can relate to performing a medical (e.g., surgical) procedure on the subject. In further implementations, the environment can include a human user that interacts with an agent in the form of a user equipment, e.g., a digital assistant. The piece of user equipment provides a user interface between the user and a computer system (the same computer system that implements the reinforcement learning system or a different computer system). The user interface can allow the user to enter data to and / or receive data from the computer system, and the agent is controlled by the action selection policy to perform an information transfer task related to the user, such as providing the user with information about a topic and / or allowing the user to specify a component of a task for the computer system to perform. For example, the information transfer task can be to teach the user a skill, such as how to speak a language or how to navigate around a certain geographic location; or the task can be to allow the user to define a three-dimensional shape to the computer system, e.g., so that the computer system can control an additive manufacturing (3D printing) system to produce an object having that shape. The actions can include outputting information to the user (e.g., in a certain format, at a certain rate, etc.) and / or configuring the interface to receive input from the user. For example, the actions can include setting questions for the user to perform that are related to the skill (e.g., requiring the user to choose between multiple options to use the language correctly, or requiring the user to speak out loud a passage of the language), and / or receiving input from the user (e.g., registering a selection of one of the options, or recording a spoken passage of the language using a microphone). The reward can be generated based on a measure of the effectiveness of the task performance. For example, this can be done by measuring the effectiveness of the user learning the topic - e.g., performing instances of the skill - (e.g., as measured by an automated skill assessment unit of the computer system). In this way, a personalized teaching system can be provided that is tailored to the user's abilities and current knowledge. In another example, when the information transfer task is to specify a component of a task for the computer system to perform, the actions can include: giving the user (visually, haptically, or auditorily) a user interface that permits the user to specify an element of the component of the task; and receiving user input using the user interface.Rewards can be generated based on a measure of how well and / or how easily a user specifies components of a task for the computer system to perform (e.g., how complete or well specified a three-dimensional object is). This can be determined automatically, or the rewards can be specified by the user, e.g., a subjective measure of the user’s experience. In this way, a personalized system can be provided for the user to control the computer system, which is also tailored to the user’s abilities and current knowledge.

[0183] Optionally, in any of the above implementations, the observation at any given time step can include data from previous time steps that can be beneficial to characterize the environment, e.g., actions performed at previous time steps, rewards received at previous time steps, or both.

[0184] Figure 4 is a flowchart of an example process 400 for training a policy neural network and a value neural network. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system (e.g., the training system 100 of FIG. 1), appropriately programmed, can perform the process 400. Figure 1

[0185] The system can repeatedly perform iterations of the process 400 on different imagined trajectories generated by using the world model neural network to update the parameters of the policy neural network and the value neural network.

[0186] The policy neural network is configured to process the policy network input including the updated latent representation according to current values of the policy network parameters to generate an action selection output. The value neural network is configured to process the value network input including the updated latent representation according to current values of the value network parameters to generate a predicted state value.

[0187] The system can continue to perform iterations of the process 400 until termination criteria for the training of the neural networks have been satisfied, e.g., until the parameters have converged, until a threshold amount of wall-clock time has passed, or until a threshold number of iterations of the process 400 have been performed.

[0188] The system receives a latent representation, which is in a latent space and characterizes a current state of the environment (step 402). The latent representation can be generated by a world model neural network based on a current observation (e.g., an image or video frame) of the current state of the environment. The world model neural network can 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 can be generated by the encoder neural network based on a predetermined initial hidden state of the sequence model neural network and based on the current observation.​

[0189] The system generates an imagined trajectory that includes, at each time step in the imagined trajectory, an updated latent representation (step 404). The imagined trajectory starts with the received latent representation. That is, at the first time step, the updated latent representation is the received latent representation. In brief, as described below with reference to Figure 5 In more detail, generating each imagined trajectory can include repeatedly generating action selection outputs by using the policy neural network, selecting an action to be performed by the agent using the action selection outputs, and then generating an updated latent representation using the world model neural network based on at least the selected action.

[0190] For each updated latent representation in the imagined trajectory, the system determines a predicted reward (step 406) and determines 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 the cumulative measure of the predicted reward that the agent would receive by using the action selection outputs to select actions. For example, the predicted state value can specify the sum of long-term, time-adjusted (e.g., time-discounted) predicted rewards that the agent would receive by using the action selection outputs to select actions, starting from the state characterized by the updated latent representation.

[0191] In more detail, for each updated latent representation in the imagined trajectory, the system can determine the predicted reward by processing the updated latent representation and the updated hidden state of the sequence model neural network according to current values of the reward network parameters using the reward neural network to generate the predicted reward, which can involve using the sequence model neural network to process the current latent representation and the selected action to generate the updated hidden state. Likewise, the system can determine the predicted state value by processing the updated latent representation using the value neural network to generate the predicted state value.

[0192] The system determines, based on the predicted rewards and the predicted state values of the updated latent representations in the imagined trajectory, a corresponding target state value for each updated latent representation in the trajectory (step 410). For example, the target state value can be determined by computing a bootstrap reward that integrates the predicted rewards and the predicted state values. The above with reference to Figure 2B An example method of computing the bootstrap reward is described.

[0193] The system determines updates to the current values of the policy network parameters by optimizing the actor objective function (step 412). In some implementations, the actor objective function includes a first term that depends on the target state values. In some implementations, the actor objective function includes a second term that depends on an entropy of the action selection output generated by the policy neural network. See Figure 2B Examples of actor objective functions are described.

[0194] The system can determine the updates by, for example, computing respective gradients of the actor objective function with respect to the policy network parameters by backpropagation, and then applying an update rule, e.g., an Adam update rule, an AdamW update rule, an Adagrad update rule, or an RMSProp update rule, to the respective gradients.

[0195] The system applies a symmetric log transform to each target state value to generate a transformed target state value for each updated latent representation in the imagined trajectory (step 414).

[0196] The system encodes each transformed target state value to generate an encoded transformed target state value for each updated latent representation imagined in the trajectory (step 416). For example, the system can generate the encoded transformed target state values using a duohot encoding technique, which can each be represented by a vector of a predetermined length having a predetermined number of non-zero entries.

[0197] The system determines updates to the current values of the value network parameters by optimizing the critic objective function (step 418). The critic objective function includes a term that measures a difference between the predicted state values and the encoded transformed target state values. See Figure 2B Examples of critic objective functions are described.

[0198] The system can determine the updates by, for example, computing respective gradients of the critic objective function with respect to the value network parameters by backpropagation, and then applying an update rule, e.g., an Adam update rule, an AdamW update rule, an Adagrad update rule, or an RMSProp update rule, to the respective gradients.

[0199] Figure 5 is a flowchart of an example process 500 for generating an imagined trajectory of latent representations. For convenience, the process 500 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system (e.g., the training system 100 of FIG. 1) appropriately programmed in accordance with this specification can perform the process 500. Figure 1 The training system 100 of FIG. 1 can perform the process 500.

[0200] The system can repeatedly perform the process 500 to generate an imagined trajectory that includes an updated latent representation at each of a plurality of time steps during the imagined trajectory. The imagined trajectory starts with the latent representation received at step 402 of the process 400. That is, at the first time step, the updated latent representation is the received latent representation.

[0201] The system processes 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 (step 502). For the first time step, the current latent representation is the received latent representation. For any subsequent time steps, the current latent representation is the updated latent representation that has been generated in the previous time step.

[0202] The system uses the action selection output to select an action that the agent will perform when the environment is in the state characterized by the current latent representation (step 504). The above references to the process 400 describe the selection of the action by using the action selection output. Figure 1 Similarly, the selection of the action by using the action selection output is described.

[0203] The system generates an updated latent representation based on the current latent representation and the selected action (step 506). The updated latent representation characterizes a state to which the environment will transition if the agent performs the selected action when the environment is in the state characterized by the current latent representation.

[0204] To this end, the system processes the current latent representation and the selected action using the sequence model neural network and in accordance with current values of the sequence model network parameters to generate an updated hidden state. The system then processes the updated hidden state using the dynamics predictor neural network to generate the updated latent representation.

[0205] Before training the policy neural network and the value neural network by performing the processes 400 and 500, the system trains the world model neural network to ensure the quality of the imagined trajectory. Specifically, the system trains the encoder neural network and the dynamics predictor neural network to encourage the generation of high-quality latent representations, e.g., latent representations that can capture significant (e.g., task-specific) information contained in observations of the environment, or latent representations that predict subsequent states of the environment. The system also trains the sequence model neural network to more accurately model the potential dynamics of the environment, i.e., the state transitions of the environment that will result from the agent performing proposed actions; trains the reward neural network to more accurately estimate the rewards that the agent will receive from the environment in response to performing selected actions; and trains the continuation predictor neural network to more accurately predict whether a current state is a terminal state of the environment.

[0206] Figure 6is a flowchart of an example process 600 for training a world model neural network. For convenience, the process 600 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system (e.g., the training system 100 of Figure 1 FIG. 1) appropriately programmed in accordance with this specification, can perform process 600.

[0207] The system can sample historical trajectories from a replay buffer and train the world model neural network on the sampled historical trajectories. The replay buffer can be any available memory that stores historical trajectories that were generated as a result of an agent (or another agent) interacting with an environment (or with another instance of the environment) to attempt to perform one or more tasks.

[0208] In some implementations, each historical trajectory can include a tuple at each of a plurality of time points between the historical trajectory. At a given time point, the tuple includes: (i) a current observation that characterizes a current state of the environment, (ii) a current action performed by the agent, (iii) a current reward received by the agent, and (iv) a true value continuation flag that indicates whether the current state of the environment is a terminal state of the environment (e.g., a goal state of the environment that is achieved after completing a task).

[0209] The system can repeatedly perform iterations of the process 600 on different historical trajectories sampled from the replay buffer to update parameters of the world model neural network. The system can continue to perform iterations of the process 600 until a termination criterion for training of the world model neural network has been satisfied, e.g., until the parameters have converged, until a threshold amount of wall-clock time has elapsed, or until a threshold number of iterations of the process 600 have been performed.

[0210] In general, the system can perform the following steps 602-612 at each of a plurality of time steps during a historical trajectory.

[0211] The system processes (i) a current latent representation that was generated at a previous time step and that characterizes a previous state of the environment, and (ii) data that specifies a previously selected action that was selected at the previous time step using the sequence model neural network in accordance with current values of the sequence model network parameters, to generate an updated hidden state that characterizes the current state of the environment (step 602).

[0212] The system processes the updated hidden state and a current observation of the current state of the environment using the encoder neural network in accordance with current values of the encoder network parameters to generate an updated latent representation that characterizes the current state of the environment (step 604).

[0213] The system processes the updated hidden state using the dynamic predictor neural network in accordance with current values of the dynamic predictor network parameters to generate an estimated updated latent representation (step 606). The estimated updated latent representation is an estimated value of the updated latent representation that characterizes the current state of the environment and that can be generated by the encoder neural network.

[0214] The system processes the updated latent representation and optionally the updated hidden state using the decoder neural network in accordance with current values of the decoder network parameters to generate a reconstructed current observation of the current state of the environment (step 608).

[0215] The system processes the updated latent representation and optionally the updated hidden state using the reward neural network in accordance with current values of the reward network parameters to generate a predicted reward (step 610). The predicted reward is a prediction of a reward that the agent will receive when the environment is in the current state.

[0216] The system processes the updated latent representation and optionally the updated hidden state using the continuation predictor neural network in accordance with current values of the decoder network parameters to generate a predicted continuation predictor flag that indicates whether the current state is a terminal state of a trajectory (step 612).

[0217] The system determines an update to the current values of the parameters of the world model neural network based on optimizing a combined objective function (step 614). Generally, the combined objective function measures a discrepancy between outputs of the neural network and true value information included in the historical trajectories. The combined objective function is described above with reference to Figure 2A Examples of the combined objective function are described.

[0218] More specifically, determining the update to the current values of the parameters of the world model neural network can involve determining an update to current values of the sequence model neural network parameters and the dynamic predictor neural network parameters based on computing gradients of the dynamic objective function with respect to the sequence model neural network parameters and the dynamic predictor neural network parameters. The dynamic objective function depends on a Kullback-Leibler divergence between (i) a latent representation generated by the encoder neural network and (ii) a latent representation generated by the dynamic predictor neural network. Here, the sequence model neural network parameters are parameters of a sequence model neural network configured to process the current latent representation and the selected action to generate an updated hidden state of the sequence model neural network. The dynamic predictor neural network parameters are parameters of a dynamic predictor neural network configured to process the updated hidden state to generate the updated latent representation. The encoder neural network is configured to process the observation to generate the latent representation that characterizes the current state of the environment.

[0219] The system can determine the update by, for example, computing a respective gradient of the combined objective function with respect to appropriate network parameters by backpropagation, and then applying an update rule, e.g., an Adam update rule, an AdamW update rule, an Adagrad update rule, or an RMSProp update rule, to the respective gradient.

[0220] During training, the system can incorporate any number of techniques to improve the speed, effectiveness, or both, of the training process. For example, the system can avoid using hyperparameter annealing, weight decay, or dropout. As another example, the system can initialize the values of the parameters of some of the above-described neural networks, e.g., the reward network parameters and the value network parameters, to zero. This can speed up early learning by mitigating a problem of latency at the start of training caused by a large number of predicted rewards, which can otherwise occur if the parameters are randomly initialized at the start of training.

[0221] It will be appreciated that the techniques for training an encoder-decoder neural network can be applied to any of a wide variety of tasks, e.g., tasks that can not involve controlling an agent. As a general example, an encoder-decoder neural network can 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.

[0222] An encoder-decoder neural network can operate on any of a variety of different types of data items, e.g., images, videos, audio signals, etc. As a particular example, an encoder-decoder neural network can operate on high-dimensional (and thus large) data items. For example, an encoder-decoder neural network can be used to compress and reconstruct high-resolution images, e.g., medical images and geographic images. In more detail, an input data item can comprise an image of a real-world entity captured by an image sensor, such as a camera or a LIDAR sensor, or captured by a medical image sensor, such as an X-ray or MRI machine. The latent representation of the input data item constitutes a compressed version of the input data item, e.g., a compressed version that can be stored and / or transmitted over a communication network. The compressed version of the input data item can be retrieved from storage and / or received over a communication network, and the input data item reconstructed.

[0223] Figure 7 FIG. 7 is a flow diagram of an example process 700 for training an encoder-decoder neural network. For convenience, the process 700 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 of FIG. 1, appropriately programmed, can perform the process 700. Figure 1 The training system 100 of FIG. 1, or another training system that includes a symmetrical log transform engine, can perform the process 700.

[0224] The system receives a training data item (step 702).

[0225] The system processes the training data item using an encoder-decoder neural network to determine a latent representation of the training data item and generates a reconstruction of the training data item from current values of parameters of the encoder-decoder neural network (step 704). The latent representation is a lower-dimensional, i.e., more compact, version of the training data item. The encoder-decoder neural network includes an encoder neural network and a decoder neural network.

[0226] More specifically, the system processes the training data item using the encoder neural network to determine the latent representation of the training data item. The system processes the latent representation of the training data item using the decoder neural network to generate an output of the encoder-decoder neural network. The system then applies a symmetric exponential transform to the output of the encoder-decoder neural network to generate the reconstruction of the training data item. For example, the symmetric exponential transform can take the following form: or where x is the output of the encoder-decoder neural network.

[0227] The system applies a symmetric logarithmic transform to the training data item to generate a transformed training data item (step 706). For example, the symmetric logarithmic transform takes the following form: or where x is the training data item.

[0228] The system determines an update to the current values of the encoder-decoder network parameters based on evaluating a loss function that measures a 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 update by, for example, computing respective gradients of a combined objective function with respect to the parameters of the encoder-decoder neural network by backpropagation, and then applying an update rule, e.g., an Adam update rule, an AdamW update rule, an Adagrad update rule, or an RMSProp update rule, to the respective gradients.

[0229] Figure 8 Quantitative examples of performance gains that can be achieved across a large number of tasks spanning 7 different domains, including continuous and discrete actions, visual and low-dimensional inputs, dense and sparse rewards, different scales of rewards, 2D and 3D worlds, and procedural generation, are shown by using the neural networks described in this specification.

[0230] These tasks include proprioceptive-sensor-based control tasks and vision control tasks (described in arXiv:1801.00690), Atari tasks (described in arXiv:1903.00374 and arXiv:1207.4708), Behavior Suite tasks or bsuite tasks (described in arXiv:1908.03568), and Crafter tasks (described in arXiv:2109.06780).

[0231] Figure 8 The vertical axis in FIG. 1 shows the returns received by the agent when controlled using each of the plurality of neural network systems, where higher returns indicate better performance. In Figure 8 In FIG. 1, the DreamerV3 system corresponds to a system including the neural networks described in this specification. The other systems are existing reinforcement learning systems for agent control, e.g., the MPO system is described in arXiv:1806.06920, the DDPG system is described in arXiv:1509.02971, and the D4PG system is described in arXiv:1804.08617. It will be appreciated that DreamerV3 outperforms these existing systems on all of these tasks. Notably, the DreamerV3 system can achieve this performance gain across these tasks after being trained using the same values for this set of hyperparameters as listed in Table 1 above.

[0232] The specification uses the term“configured” when referring to systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

[0233] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can 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 to control the operation of, data processing apparatus. The computer storage medium can be machine-readable storage, machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or additionally, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0234] 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. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0235] A computer program, which can also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages, and it 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, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and are interconnected by a data communication network.

[0236] ​In this specification, the term“database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and can be stored on storage devices in one or more locations. Thus, for example, an index database can include multiple collections of data, each of which can be organized differently and accessed in different ways.

[0237] Similarly, in this specification, the term“engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components that are installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer(s).

[0238] The processes and logic flows described in this specification can be performed 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 can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0239] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read-only memory or 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 the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto-optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, 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.

[0240] 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, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0241] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device of the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0242] Data processing devices for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing the general and compute-intensive parts of machine learning training or production (i.e., inference, workloads).

[0243] Machine learning models can be implemented and deployed using a machine learning framework, such as a TensorFlow framework or a JAX framework.

[0244] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back-end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that 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 implementation of the subject matter described in this specification, or 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) and a wide area network (WAN), e.g., the Internet.

[0245] The computing system can include clients and servers. A client and server 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 transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at a server from the device.

[0246] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any inventions, but rather as descriptions of features that can be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features can be described above as acting in certain combinations and even initially claimed as such, one or more features from a claimed combination can in some cases be excised from the combination and the claimed combination can be directed to a subcombination or variation of a subcombination.

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

[0248] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the acts recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve the desired results. In some instances, multitasking and parallel processing can be advantageous.

Claims

1. A method of training a policy neural network having a plurality of policy network parameters for selecting actions to be performed by an agent in an environment, the method comprising: receiving a latent representation, the latent representation being in a latent space and characterizing a current state of the environment; generating an imagined trajectory of latent representations, starting from the received latent representation, by performing the following at each of a plurality of time steps: processing a current latent representation using the policy neural network to generate an action selection output in accordance with current values of the policy network parameters; using the action selection output to select an action to 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 characterizing 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; for each latent representation in the imagined trajectory: determining a predicted reward; and processing the latent representation using a value neural network having a plurality of 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 using action selection outputs to select actions starting from the state characterized by the latent representation in accordance with current values of the value network parameters; determining a target state value for each latent representation in the imagined trajectory as a function of the predicted reward and the predicted state value for the latent representation in the imagined trajectory; determining an update to the current values of the policy network parameters by optimizing an actor objective function comprising a first term that depends on the target state value; applying a symmetric logarithmic transform to each target state value to generate a transformed target state value for each latent representation in the imagined trajectory; encoding each transformed target state value to generate an encoded transformed target state value for each latent representation in the imagined trajectory; and determining an update to the current values of the value network parameters by optimizing a critic objective function that measures a difference between the predicted state value and the encoded transformed target state value.

2. The method of claim 1, wherein the symmetric log transform has the form: or where x is the target state value.

3. The method of any one of claims 1-2, wherein the encoded transformed target state value is represented by a vector of a predetermined length having a predetermined number of non-zero entries.

4. The method of any one of claims 1-3, wherein generating the predicted state value for each latent representation in the imagined trajectory comprises: processing the latent representation to assign a respective probability to each predetermined interval in a sequence of predetermined intervals between possible state values; generating an output of the value neural network by computing a weighted average of the state values represented by the sequence of predetermined intervals weighted by the respective probabilities; and applying a symmetric exponential transform to the output of the value neural network to generate the predicted state value for the latent representation in the imagined trajectory.

5. The method of claim 4, wherein the symmetric exponentiation transformation has the form: or where x is the output of the value neural network.

6. The method of any one of claims 1-5, wherein determining the update to the current values of the policy network parameters by optimizing the actor objective function comprises: generating a scaled-down target state value for each latent representation in the imagined trajectory from the target state value using a scaling factor; and determining the update to the current values of the policy network parameters by optimizing the actor objective function including the first term that depends on the scaled-down target state value.

7. The method of claim 6, wherein generating the scaled-down target state value comprises: dividing the target state value by a maximum of a given integer value and the scaling factor, wherein the scaling factor is computed using a respective percentile of one or more target state values of the latent representations in the imagined trajectory.

8. The method of any one of claims 1-7, wherein the actor objective function includes a second term that depends on a scaled entropy of the action selection output generated by the policy neural network, wherein the scaled entropy is scaled by a predetermined entropy scaling factor.

9. The method of any one of claims 1-8, wherein processing the current latent representation and the selected action to generate the updated latent representation comprises: processing the current latent representation and the selected action using a sequence model neural network to update a hidden state of the sequence model neural network according to current values of sequence model network parameters 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 to generate the updated latent representation according to current values of dynamics predictor network parameters.

10. The method of claim 9, wherein receiving the latent representation that characterizes the current state of the environment comprises: processing an observation that characterizes the current state of the environment using an encoder neural network to generate the latent representation according to current values of encoder network parameters.

11. The method of claim 10, further comprising: processing the latent representation and the updated hidden state using a decoder neural network to generate a reconstruction of the observation according to current values of decoder network parameters; applying the symmetric log transform to the observation to generate a transformed observation; and determining an update to the current values of the decoder network parameters based on computing a gradient of a prediction objective function that includes a first term that measures a difference between the reconstruction of the observation and the transformed observation.

12. The method of any one of claims 1-11, wherein determining the predicted reward for each latent representation in the imagined trajectory comprises: processing the latent representation and the updated hidden state of the sequence model neural network using a reward neural network to generate the predicted reward according to current values of reward network parameters.

13. The method of claim 12, further comprising: applying the symmetric log transform to a reward received by the agent when the environment is in the state characterized by the latent representation to generate a transformed reward; and determining, with respect to the reward network parameters, updates to the current values of the reward network parameters based on computing gradients of a predicted objective function, the predicted objective function including a second term that measures a difference between the predicted reward and the transformed reward.

14. The method of any one of claims 1-13, further comprising: determining, based on computing gradients of the dynamic objective function with respect to the sequence model neural network parameters and the dynamic predictor neural network parameters, updates to the current values of the sequence model neural network parameters and the dynamic predictor neural network parameters, wherein the dynamic objective function depends on a Kullback-Leibler divergence between (i) a latent representation generated by an encoder neural network and (ii) a latent representation generated by the dynamic predictor neural network.

15. The method of any one of claims 1-14, wherein: 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 dynamic predictor neural network and the reward neural network are each configured as a multi-layer perceptron (MLP); and the sequence model neural network is configured as a recurrent neural network.

16. The method of claim 15, wherein the convolutional neural network, the multi-layer perceptron (MLP), or the recurrent neural network comprises: one or more layer normalization layers and one or more sigmoid linear unit (SiLU) activation layers.

17. The method of any one of claims 15-16, further comprising initializing the current values of the reward network parameters and the value network parameters to zero.

18. The method of any one of claims 1-17, further comprising using the policy neural network to select an action to be performed by the agent in the environment by processing, using the policy neural network, an input derived from data characterizing a current state of the environment at a time step to generate an action selection output for determining an action to be performed by the agent at the time step.

19. A method performed by one or more computers for training an encoder-decoder neural network, wherein the encoder-decoder neural network is 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, and wherein the method comprises: 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 log transform to the training data item to generate a transformed training data item; and determining an update to the current value of the encoder-decoder network parameters based on evaluating a loss function that measures a difference between the reconstruction of the training data item and the transformed training data item.

20. The method of claim 18, wherein the symmetric log transform has the form: or where x is the training data item.

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

22. The method of any of claims 18-20, wherein generating the reconstruction of the input data item comprises: processing the training data item to determine the latent representation of the training data item; processing the latent representation of the training data item to generate an output of the encoder-decoder neural network; and applying a symmetric exponential transform to the output of the encoder-decoder neural network to generate the reconstruction of the training data item.

23. The method of claim 21, wherein the symmetric exponentiation transform has the form: or where x is the output of the encoder-decoder neural network.

24. The method of any of claims 18-22, wherein the input data item comprises an image.

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

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 operations of the respective method of any preceding claim.

Images