Attention Neural Networks with Short-Term Memory Units

Abstract

A system for controlling an agent that interacts with an environment to perform a task, the system including an action selection neural network configured to generate an action selection output used to select an action to be performed by the agent, the action selection neural network including an encoder sub-network configured to generate an encoded representation of a current observation, an attention sub-network configured to generate an attention sub-network output using an attention mechanism, a recurrent sub-network configured to generate a recurrent sub-network output, and an action selection sub-network configured to generate an action selection output used to select an action to be performed by the agent in response to the current observation.

Description

[Technical field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Patent Application No. 63 / 146,361, filed February 5, 2021. The disclosure of the prior application is deemed to be part of and incorporated by reference into the disclosure of this application.

[0002] This specification relates to reinforcement learning. [Background technology]

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

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

[0005] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output for a received input. Some neural networks are deep neural networks that include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input for the next layer in the network, i.e., the next hidden layer, or the output layer. Each layer of the network generates an output from the received input depending on the current values ​​of the respective set of parameters. [Prior art documents] [Non-patent literature]

[0006] [Non-Patent Document 1] Vaswani et al., Attention Is All You Need, arXiv:1706.03762 [Non-Patent Document 2] Parisotto et al., Stabilizing transformers for reinforcement learning, arXiv:1910.06764 [Non-Patent Document 3] Song et al., V-mpo: On-policy maximum a posteriori policy optimization for discrete and continuous control, arXiv:1909.12238 [Non-Patent Document 4] Kapturowski et al., Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2018 [Non-Patent Document 5] DeepMind Lab(https: / / arxiv.org / abs / 1612.03801) Summary of the Invention [Means for solving the problem]

[0007] This document generally describes a reinforcement learning system that controls an agent that interacts with an environment.

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

[0009] By incorporating the self-attention mechanism of an attention-based neural network together with the memory mechanism of a recurrent neural network, such as a long short-term memory (LSTM) neural network, into the action selection neural network used by a reinforcement learning system to select actions to be performed by an agent, the described technique can provide temporally structured information to the action selection neural network to improve the quality of the action selection output during or after training, i.e., at run-time. In particular, the described technique effectively leverages both the self-attention mechanism that extracts long-range dependencies and the memory mechanism that reasons about short-term dependencies to integrate information about the agent's past interactions with the environment at multiple different time scales, thereby enabling the action selection neural network to reason about events across multiple time scales and adjust future action selection policies accordingly.

[0010] Furthermore, the techniques described herein optionally include the implementation of a trainable gating mechanism to enable the action selection neural network to more effectively combine information calculated by using the attention-based neural network and the recurrent neural network. This effective combination may be particularly advantageous in complex environmental settings, since it allows greater flexibility in determining which information should be processed in controlling the robotic agent. Here, the term "gating mechanism" refers to a unit that forms a dataset based on both the input to the neural network and the output of the neural network. A trainable gating mechanism is a gating mechanism in which the dataset is further based on one or more adjustable parameter values. The gating mechanism may be used, for example, to generate a dataset based on both the input to the attention sub-network and the output of the attention sub-network.

[0011] Thus, the reinforcement learning system described herein can achieve superior performance to conventional reinforcement learning systems in controlling an agent to perform a task, for example, by receiving more cumulative extrinsic rewards. The reinforcement learning system described herein trains the action selection neural network faster than conventional reinforcement learning systems that do not utilize a self-attention mechanism or a memory mechanism, or neither of them. Furthermore, by training the action selection neural network on a control learning assistance task in addition to training the neural network to maximize cumulative rewards, the reinforcement learning system described herein can augment the feedback signal received during the training of the action selection neural network to further improve the training, for example, to encourage the learning of representations that assist with obstacle avoidance or trajectory planning. Thus, the reinforcement learning system described herein allows for more efficient use of computational resources in training.

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

[0013] [Figure 1] FIG. 1 illustrates an exemplary reinforcement learning system. [Diagram 2] 4 is a flowchart illustrating an example process for controlling an agent. [Diagram 3] 1 is a flowchart illustrating an exemplary process for determining updates to parameter values ​​of an attention selection neural network. [Figure 4] FIG. 1 illustrates determining updates to parameter values ​​of an attention selection neural network. [Diagram 5]FIG. 1 illustrates a quantitative example of the performance gains that may be realized by using the control neural network system described herein. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

[0014] Like reference numbers and designations in the various drawings refer to like elements.

[0015] This specification describes a reinforcement learning system that controls an agent that interacts with an environment by processing data (i.e., "observations") that characterize the current state of the environment at each time step among multiple time steps to select an action to be taken by the agent at that time step.

[0016] At each time step, the state of the environment at that time step depends on the state of the environment at the previous time step and on the actions taken by the agent at the previous time step.

[0017] In some implementations, the environment is a real-world environment and the agent is a machine agent that interacts with the real-world environment, e.g., a robot that moves through the environment, or an autonomous or semi-autonomous ground, air, or sea vehicle.

[0018] In these implementations, the observations may include, for example, one or more of images capturing the observations as the agent interacts with the environment, object position data, and sensor data, e.g., from image, distance, or position sensors, or from actuators.

[0019] For example, in the case of a robot, the observations may include data characterizing the current state of the robot, such as one or more of joint positions, joint velocities, joint forces, torques or accelerations, such as gravity compensated torque feedback, global or relative poses of items being held by the robot.

[0020] In the case of a robot or other mechanical agent, or vehicle, the observations may similarly include one or more of the position, linear or angular velocity, force, torque or acceleration, and global or relative pose of one or more parts of the agent. The observations may be defined in one, two, or three dimensions and may be absolute and / or relative observations.

[0021] The observations may also include sensed electronic signals, such as, for example, motor current or temperature signals, and / or image data or video data, for example, from a camera or LIDAR sensor, such as data from a sensor on the agent or from a sensor located separately from the agent in the environment.

[0022] In these implementations, the actions may be control inputs that control a robot, e.g., torques on the robot's joints or high-level control commands, or control inputs that control an autonomous or semi-autonomous ground, air, or sea vehicle, e.g., torques on the vehicle's control surfaces or other control elements, or high-level control commands.

[0023] In other words, an action may include, for example, position data, velocity data, or force / torque / acceleration data for one or more joints of a robot, or one or more parts of another mechanical agent. Action data may also or alternatively include electronic control data, such as motor control data, or more generally, data for controlling one or more electronic devices in the environment, where control affects the observed state of the environment. For example, in the case of an autonomous or semi-autonomous ground, air, or sea vehicle, actions may include actions that control the navigation, e.g., steering, and movement, e.g., braking and / or acceleration, of the vehicle.

[0024] In some other applications, the agent may control actions in a real-world environment that includes items of equipment, e.g., in a data center, in a power / water distribution system, or in a manufacturing plant or service facility. The observations may then relate to the operation of the plant or facility. For example, the observations may include observations of power or water usage by the facility, or may include observations of power generation or distribution control, or may include observations of resource usage or waste production. The actions may include actions that control or impose operating conditions on the items of equipment of the plant / facility, and / or actions that result in a change in settings in the operation of the plant / facility, e.g., to adjust or turn on / off components of the plant / facility.

[0025] In the case of electronic agents, the observations may include data from one or more sensors monitoring parts of a plant or service facility, such as current sensors, voltage sensors, power sensors, temperature sensors, and other sensors, and / or electronic signals representing the function of electronic and / or mechanical items of equipment. For example, the real-world environment may be a manufacturing plant or service facility, and the observations may relate to the operation of the plant or facility, e.g., resource usage such as power consumption, and the agent may control actions or operations at the plant / facility, e.g., to reduce resource usage. In some other implementations, the real-world environment may be a renewable energy plant, and the observations may relate to the operation of the plant, e.g., to maximize current or future planned power generation, and the agent may control actions or operations at the plant to achieve this.

[0026] As another example, the environment may be a chemical synthesis environment or a protein folding environment, with each state being a respective state of a protein chain or of one or more intermediate or precursor chemicals, and the agent is a computer system for determining how to fold a protein chain or synthesize a chemical. In this example, the actions are possible folding actions for folding a protein chain or actions for assembling precursor chemicals / intermediates, and the results to be achieved may include, for example, folding a protein so that it is stable and that it achieves a particular biological function, or resulting in a plausible synthetic route for a chemical. As another example, the agent may be a machine agent that performs or controls protein folding actions or chemical synthesis steps selected by the system automatically without human interaction. The observations may include direct or indirect observations of the state of the protein or chemicals / intermediates / precursors, and / or may be derived from simulations.

[0027] In some implementations, the environment may be a simulated environment and the agent may be implemented as one or more computers that interact with the simulated environment.

[0028] The simulated environment may be a motion simulation environment, such as a driving or flight simulation, and the agent may be a simulated vehicle that moves within the motion simulation. In these implementations, the actions may be control inputs that control a simulated user or a simulated vehicle.

[0029] In some implementations, the simulated environment may be a simulation of a particular real-world environment. For example, the system may be used to select actions in the simulated environment during training or evaluation of the control neural network, and after training or evaluation, or both, are completed, may be deployed to control a real-world agent in a real-world environment simulated by the simulated environment. This may avoid unnecessary wear and tear and damage to the real-world environment or real-world agent, and may allow the control neural network to be trained and evaluated for situations that rarely occur or are difficult to replicate in real-world environments.

[0030] In general, in the case of a simulated environment, the observations may include simulated versions of the observations described above, or one or more of the types of observations described above, and the actions may include simulated versions of the actions described above, or one or more of the types of actions described above.

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

[0032] 1 illustrates an exemplary reinforcement learning system 100. Reinforcement learning system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations in which the systems, components, and techniques described below are implemented.

[0033] The system 100 controls the agent 102 interacting with the environment 104 by selecting an action 106 to be performed by the agent 102, and then causing the agent 102 to perform the selected action 106, such as by transmitting control data to the agent 102 that instructs the agent 102 to perform the action 106. In some cases, the reinforcement learning system 100 may be implemented on or be a component of the agent 102, and the control data may be transmitted to the agent's actuator(s).

[0034] Execution of a selected action 106 by the agent 102 generally transitions the environment 104 to successive new states. By repeatedly having the agent 102 operate in the environment 104, the system 100 can control the agent 102 to complete a specified task.

[0035] System 100 includes a control neural network system 110 and one or more memories that store a set of model parameters 118 (“network parameters”) for the neural network included in control neural network system 110 .

[0036] At a high level, the control neural network system 110 is configured to process inputs including a current observation 108 characterizing a current state of the environment 104, at each time step of multiple time steps, according to the model parameters 118, to generate an action selection output 122 that can be used to select a current action 106 to be performed by the agent 102 in response to the current observation 108.

[0037] The control neural network system 110 includes an action selection neural network 120. The action selection neural network 120 is implemented with a neural network architecture that improves the quality of the action selection output 122 by enabling the system to detect and react to events occurring at different time scales. Specifically, the action selection neural network 120 includes an encoder sub-network 124, an attention sub-network 128, a gating sub-network 132 (preferably), a recurrent sub-network 136, and an action selection sub-network 140. Each sub-network can be implemented as a group of one or more neural network layers in the action selection neural network 120.

[0038] At each time step of the multiple time steps, the encoder sub-network 124 receives an encoder sub-network input including a current observation 108 that characterizes the current state of the environment 104, and generates an encoded representation of the current observation 108 ("Y tThe system is configured to process the encoder sub-network inputs according to trained parameter values ​​of the encoder sub-network to generate an encoded representation 126. The encoded representation 126 can be in the form of an ordered collection of values, e.g., a vector or matrix of values. For example, the encoded representation 126, which is then provided as an input to the attention sub-network 128, can be an input vector having a respective input value at each input position of a number of input positions in an input order. In some implementations, the encoded representation 126 has the same dimensions as the observations 108, while in other implementations, the encoded representation 126 has smaller dimensions compared to the observations 108 for computational efficiency reasons. In some implementations, in addition to providing the encoded representation 126 as an input to the attention sub-network 128, the system also stores the encoded representation 126 generated at a given time step in a memory buffer or lookup table so that the encoded representation 126 can be used later, e.g., in future time steps following the given time step.

[0039] If the observations are images, the encoder sub-network 124 may be a convolutional sub-network configured to process the observations for a time step, e.g., a convolutional neural network with residual blocks. In some cases, e.g., if the observations include lower dimensional data, the encoder sub-network 124 may additionally or instead include one or more fully connected neural network layers.

[0040] The attention sub-network 128 is a network that includes one or more attention neural network layers. Each attention layer operates on a respective input sequence that includes a respective input vector at each of one or more positions (e.g., a plurality of concatenated input vectors). At each time step of the multiple time steps, the attention sub-network 128 receives an attention sub-network input that includes an encoded representation 126 of the current observation 108 and an encoded representation of one or more previous observations, and generates an attention sub-network output ("X t The attention sub-network is configured to process the attention sub-network inputs according to the trained parameter values ​​of the attention sub-network to generate the attention sub-network output ("X t " ) 130 is an output determined or otherwise derived from an updated (i.e., "attention-directed") representation of the encoded representation generated by using one or more attention neural network layers.

[0041] In particular, in addition to the encoded representation 126 of the current observation 108, the attention sub-network input also includes encoded representations of one or more previous observations characterizing one or more previous states of the environment that immediately precede the current state of the environment 108. Each encoded representation of a previous observation may be in the form of a respective input vector that includes a respective input value at each input position of a number of input positions in the input sequence. Thus, the attention sub-network input may be a concatenated input vector consisting of a number of individual input vectors, each corresponding to a respective encoded representation of an observation of a different previous state of the environment 108 leading up to (and including) the current state.

[0042] In general, the attention layer in the attention sub-network 128 can be configured in any of a variety of configurations. Examples of configurations of attention layers, as well as details of other components of attention sub-networks, are described in more detail in Vaswani et al., Attention Is All You Need, arXiv:1706.03762, and Parisotto et al., Stabilizing transformers for reinforcement learning, arXiv:1910.06764, which are incorporated by reference in their entireties. For example, the attention mechanism applied by the attention layer in the attention sub-network 128 can be a self-attention mechanism, e.g., a multi-head self-attention mechanism.

[0043] In general, the attention mechanism maps a query and a set of key-value pairs to an output, where the query, key, and value are all vectors derived from the input to the attention mechanism based on the respective matrices. The output is calculated as a weighted sum of values, where the weight assigned to each value is calculated by a fitness function, e.g., a dot product or a scaled dot product, of the query with the corresponding key. In general, the attention mechanism determines the relationship between two sequences, while the self-attention mechanism is configured to relate different positions in the same sequence and determine a transformed version of the sequence as the output. The attention layer input may include vectors for each element of the input sequence. These vectors provide input to the self-attention mechanism and are used by the self-attention mechanism to determine a new representation of the same sequence for the attention layer output, which also includes vectors for each element of the input sequence. The output of the self-attention mechanism may be used as the attention layer output.

[0044] In some implementations, the attention mechanism may be, for example, Q A query transformation defined by, for example, the matrix W K and, for example, the matrix W VFor each input vector X in the input sequence, we apply the value transformations defined by Q , key vector K=XW K , and the value vector V=XW V For example, the attention mechanism may be a dot-product attention mechanism that is applied by applying each query vector to each key vector to determine a respective weight for each value vector, and then combining the value vectors using the respective weights to determine the attention layer output for each element of the input sequence. The attention layer output may be scaled by a scaling factor, e.g., the square root of the dimensions of the query and keys, to implement scaled dot-product attention. Thus, for example, the output of the attention mechanism may be

[0045]

number

[0046] where d is the dimension of the key (and value) vector. In another implementation, the attention mechanism comprises an "additive attention" mechanism that uses a feedforward network with a hidden layer to compute the fitness function.

[0047] The attention mechanism may implement multi-headed attention, i.e., it may apply multiple different attention mechanisms in parallel, whose outputs may then be combined, e.g., concatenated, if necessary, with a learned linear transformation applied to reduce them back to the original dimensionality.

[0048] An attention neural network layer or a self-attention neural network layer is a neural network layer that includes an attention or self-attention mechanism that operates on attention layer inputs to generate attention layer outputs. The attention sub-network 128 may comprise a single attention layer, or alternatively, a sequence of attention layers, where each attention layer other than the first receives as input the output from the previous attention layer in the sequence.

[0049] Furthermore, in this embodiment, the self-attention mechanism can be masked such that any given position in the input sequence does not attend to any positions following that given position in the input sequence, e.g., for subsequent positions following that given position in the input sequence, the attention weight for the subsequent positions is set to 0.

[0050] The recurrent subnetwork 136 is a network that includes one or more recurrent neural network layers, such as one or more long short-term memory (LSTM) layers, one or more gated recurrent unit (GRU) layers, one or more vanilla recurrent neural network (RNN) layers, etc. The recurrent subnetwork 136 is configured to receive and process, at each time step of the multiple time steps, a recurrent subnetwork input according to trained parameter values ​​of the recurrent subnetwork to update a current hidden state of the recurrent subnetwork corresponding to that time step and generate a recurrent subnetwork output.

[0051] In particular, the recurrent sub-network input is derived from the attention sub-network output 130. In some implementations, the action selection neural network 120 may directly provide the attention sub-network output 130 as an input to the recurrent sub-network 136. Alternatively, in some implementations, the action selection neural network 120 may utilize a gating sub-network 132 to combine the output 126 of the encoder sub-network 124 and the output 130 of the attention sub-network 128. Thus, the recurrent sub-network input is derived from the output of the gating sub-network 132 of the action selection neural network 120 ("Z t " )134.

[0052] In some of these implementations, the gating mechanism implemented by the gating sub-network 132 is a fixed, summation (or concatenation) mechanism, and the gating sub-network may include a summation (or concatenation) layer configured to receive the encoded representation of the current observation 126 and the attention sub-network output 130, and to generate an output based on the representation 126 and the output 130. For example, the gating sub-network may compute a summation (or concatenation) of i) the encoded representation of the current observation 126 and ii) the attention sub-network output 130 along a given dimension of the received layer input.

[0053] In some other of these implementations, the gating mechanism implemented by the gating sub-network 132 is a learned gating mechanism that facilitates a more effective combination of information contained in or otherwise derivable from the output 126 of the encoder sub-network 124 and the output 130 of the attention sub-network 128. In these implementations, the gating sub-network 132 may include a gated recurrent unit (GRU) layer configured, i.e., configured through training, to apply a learned GRU gating mechanism to i) the encoded representation 126 of the current observation and ii) the attention sub-network output 130 according to trained parameter values ​​of the GRU layer to generate a gated output, i.e., the gating sub-network output 134, which is then provided as an input to the recurrent sub-network 136. The use of the GRU gating mechanism to combine the encoded representation 126 and the attention sub-network output 130 is further described below with reference to FIG. 2. In some implementations, a skip ("residual") connection may be configured between the encoder sub-network 124 and the gating sub-network 132, which may be configured to receive the encoded representation 126 via the skip connection in addition to receiving the attention sub-network output 130 directly from the attention sub-network 128.

[0054] An action selection sub-network 140 is configured to receive the action selection sub-network input at each time step of the multiple time steps and process the action selection sub-network input according to trained parameter values ​​of the action selection sub-network 140 to generate an action selection output 122. The action selection sub-network input includes the recurrent sub-network output, and in some implementations further includes the encoded representation 126 generated by the encoder sub-network 124. In implementations in which the action selection sub-network input also includes the encoded representation 126, a skip connection may be configured between the encoder sub-network 124 and the action selection sub-network 140, and the action selection sub-network 140 may be configured to receive the encoded representation 126 via the skip connection in addition to receiving the recurrent sub-network output directly from the recurrent sub-network 136.

[0055] The system 100 then uses the action selection output to select an action to be performed by the agent in the current time step. Several examples of using the action selection output to select an action to be performed by the agent are described next.

[0056] In one embodiment, the action selection output 122 may include a respective numerical probability value for each action in a set of possible actions that may be performed by the agent. The system may select an action to be performed by the agent, for example, by sampling the actions according to the probability values ​​for the actions, or by selecting the action with the highest probability value.

[0057] In another embodiment, the action selection output 122 may directly define the action to be performed by the agent, for example by defining the torque values ​​to be applied to the joints of the robotic agent.

[0058] In another embodiment, the action selection output 122 may include a respective Q value for each action in a set of possible actions that may be performed by the agent. The system may process the Q values ​​(e.g., using a softmax function) to generate a respective probability value for each possible action, which may be used to select the action to be performed by the agent (as described above). The system may also select the action with the highest Q value as the action to be performed by the agent.

[0059] The Q-value for an action is an estimate of the “return” that results from the agent performing the action in response to a current observation, and then selecting future actions to be performed by the agent depending on the current values ​​of the policy neural network parameters.

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

[0061] In some cases, the system may select an action to be performed by an agent according to a search policy. For example, the search policy may be a ∈-greedy search policy, where the system selects an action to be performed by an agent according to the action selection output 122 with probability 1-∈ and randomly selects an action with probability ∈. In this example, ∈ is a scalar value ranging from 0 to 1.

[0062] Controlling agents using the control neural network system 110 is described in more detail below with reference to FIG.

[0063] To more effectively control the agent 102 interacting with the environment 104, the reinforcement learning system 100 can use the training engine 150 to train the action selection neural network 120 to determine trained values ​​for the parameters 118 of the action selection neural network 120.

[0064] The training engine 150 is configured to train the action selection neural network 120 by iteratively updating the model parameters 118, i.e., the parameter values ​​of the encoder subnetwork 124, the attention subnetwork 128, the gating subnetwork 132, the recurrent subnetwork 136, and the action selection subnetwork 140, based on the interaction of the agent 102 (or another agent) with the environment 108 (or another instance of an environment).

[0065] In particular, the training engine 150 trains the action selection neural network 120 through reinforcement learning and, optionally, through contrastive representation learning. By contrastive representation learning, it is meant that the neural network components of the action selection neural network, in particular the attention sub-network 128, are taught to generate outputs when the action selection neural network receives a pair of similar inputs (e.g., measured by a similarity measure, e.g., a distance measure such as Euclidean distance) in succession, such that the action selection neural network generates respective outputs that are more similar to each other than the respective outputs the action selection network would generate from more distant inputs compared to the pair of similar inputs. In an example given below, the contrastive representation learning is based on training the action sub-network 128 to receive an input that is a masked form of the data generated by the encoder sub-network 124 at a given time step, and / or to generate an output that reconstructs the same data generated by the encoder sub-network 124 and / or the data generated by the encoder sub-network 124 at another time step.

[0066] In reinforcement learning, the action selection neural network 120 is trained based on the agent's interaction with the environment to optimize an appropriate reinforcement learning objective function. The architecture of the action selection neural network 120 is agnostic with respect to the choice of the exact RL training algorithm, and thus the RL training can be on-policy (e.g., one of the RL algorithms described in more detail in Song et al., V-mpo: On-policy maximum a posteriori policy optimization for discrete and continuous control, arXiv:1909.12238) or off-policy (e.g., one of the RL algorithms described in more detail in Kapturowski et al., Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2018).

[0067] As part of this RL training process, the observations 108 received by the reinforcement learning system 100 during the interaction are encoded into coded representations from which the action selection output 122 is generated. Thus, learning to generate coded representations that yield useful information is a key factor for successful RL training. To do so, the training engine 150 evaluates and uses time-domain contrastive learning objectives as proxy supervision signals for masked prediction training of the attention sub-network 128, i.e., for training the attention sub-network 128 to predict masked portions of the attention sub-network input. Such signals aim to learn self-attention consistent representations that encompass the appropriate information so that the action selection sub-network 140 effectively incorporates information previously observed (or extracted) by the attention sub-network 128.

[0068] In some implementations, by utilizing contrastive representation learning techniques to help determine parameter value updates, the training engine 150 improves the efficiency of the training process, for example, in terms of the amount of computational resources or real time consumed by the training process required to train the action selection neural network 120 to achieve or surpass state-of-the-art performance in controlling an agent to perform a given task.

[0069] Training the action selection neural network 120 is described in more detail below with reference to FIGS.

[0070] 2 is a flow chart of an exemplary process 200 for controlling an agent. For convenience, process 200 is described as being performed by one or more computer systems located at one or more locations. For example, a suitably programmed reinforcement learning system, such as reinforcement learning system 100 of FIG. 1, may perform process 200.

[0071] In general, the system may repeatedly execute process 200 at each time step of multiple time steps to select a respective action (hereinafter referred to as the “current” action) to be performed by the agent in a respective state (hereinafter referred to as the “current” state) of the environment that corresponds to that time step (hereinafter referred to as the “current” time step).

[0072] The system receives a current observation that characterizes a current state of the environment at a current time step, and generates an encoded representation of the current observation by using an encoder sub-network (step 202). For example, the current observation may include an image, a video frame, an audio data segment, a sentence in a natural language, or the like. In some of these embodiments, the observation may also include information derived from a previous time step, such as a previous action performed, a reward received in a previous time step, or both. The encoded representation of the observation may be represented as an ordered collection of numerical values, such as a vector or matrix of numerical values.

[0073] The system processes the attention sub-network input, including an encoded representation of the current observation and an encoded representation of one or more previous observations, by using the attention sub-network to generate an attention sub-network output (step 204). The one or more previous observations characterize one or more previous states of the environment that precede the current state of the environment, and thus the encoded representation of the one or more previous observations may be encoded representations generated by using the encoder sub-network at one or more time steps preceding the current time step.

[0074] The attention sub-network can be a neural network that includes one or more attention neural network layers and is configured to generate the attention sub-network output, at least in part, by applying an attention mechanism, e.g., a self-attention mechanism, to the encoded representation of a current observation and to the encoded representation of one or more previous observations characterizing one or more previous states of the environment. This use of the attention mechanism facilitates the connection of long-distance data dependencies, e.g., across respective observations of a prolonged sequence of different states of the environment.

[0075] More specifically, the attention sub-network input may include a concatenated input vector of multiple individual input vectors, each corresponding to a different encoded representation having a respective input value at each input position of the multiple input positions in the input order. To generate the attention sub-network outputs, each attention layer included in the attention sub-network may be configured to receive, for each of the one or more layer input positions, an attention layer input (which may also be in the format of a vector), and, for each particular layer input position in the layer input order, apply an attention mechanism to the attention layer input at the layer input position using one or more queries derived from the attention layer input at the particular layer input position to generate a respective attention layer output for the particular layer input position.

[0076] The system generates a combination of i) an encoded representation of the current observation and ii) an attention sub-network output, and then provides the combination as an input to the recurrent sub-network. To generate the combination, the system can utilize a gating sub-network configured to apply a gating mechanism to i) the encoded representation of the current observation and ii) the attention sub-network output to generate the recurrent sub-network input. For example, the gating sub-network can include a gated recurrent unit (GRU) layer that can be configured to apply a lower complexity GRU gating mechanism compared to an LSTM layer by utilizing a smaller number of layer parameters. As another example, the gating sub-network can compute an addition or concatenation of i) the encoded representation of the current observation and ii) the attention sub-network output.

[0077] Specifically, in the previous embodiment, the GRU layer is a recurrent neural network layer that provides similar performance to a long short-term memory (LSTM) layer with a forget gate, but has fewer parameters compared to an LSTM because it lacks an output gate. In some implementations, this gating mechanism can be applied as an update to the GRU layer that is unwound in depth of the action selection neural network instead of being unwound in time. This means that while the GRU layer is a recurrent neural network (RNN) layer, the gating mechanism can instead generate "updated" combinations of inputs received in the gating sub-network of the action selection neural network using the same formula that the GRU layer uses to update the hidden state of the GRU layer over time.

[0078] In these implementations, the GRU layer has reset gates r and update gates z, i.e.

[0079]

number

[0080] To compute each of the layers, we apply a nonlinear function, such as a sigmoid activation σ(), to a weighted combination of the received layer inputs, i.e., the encoded representation Y t and the attention subnetwork output X t Apply to the updated hidden state

[0081]

number

[0082] That is,

[0083]

number

[0084] To generate the tanh activation, we apply a nonlinear function such as tanh() to the encoded representation Y t The weighted combination of the reset gate r and the attention sub-network output X t Applying the element-wise product between

[0085]

number

[0086] are the weights (or biases) determined from the values ​​of the GRU layer parameters,

[0087]

number

[0088] represents element-wise multiplication. Then, the GRU layer produces a gated output g (l) (x,y) (which can be used as a recurrent sub-network input) as follows:

[0089]

number

[0090] Generate.

[0091] The system processes the recurrent subnetwork inputs by using the recurrent subnetwork to generate recurrent subnetwork outputs (step 206). The recurrent subnetwork can be configured to receive the recurrent subnetwork inputs and update the current hidden state of the recurrent subnetwork by processing the received inputs, i.e., to modify the current hidden state of the recurrent subnetwork generated by processing the previous recurrent subnetwork inputs by processing the currently received recurrent subnetwork inputs. The updated hidden state of the recurrent subnetwork after processing the recurrent subnetwork inputs is hereafter referred to as the hidden state corresponding to the current time step. Once the updated hidden state corresponding to the current time step is generated, the system can use the updated hidden state of the recurrent subnetwork to generate the recurrent subnetwork outputs.

[0092] For example, a recurrent sub-network can be a recurrent neural network that includes one or more long short-term memory (LSTM) layers. Due to the sequential nature of LSTM layers, they can effectively capture short-range dependencies, for example across successive observations of recent states of the environment.

[0093] The system processes the action selection sub-network inputs, including the recurrent sub-network outputs, by using the action selection sub-network to generate an action selection output that is used to select an action to be performed by the agent in response to the current observation (step 208). In some implementations, the action selection sub-network inputs also include an encoded representation. In these implementations, to generate the action selection sub-network inputs, the system can compute the concatenation of i) the recurrent sub-network outputs and ii) the encoded representation generated by the encoder sub-network at the current time step.

[0094] The system can then cause the agent to perform the selected action, i.e., by commanding the agent to perform the action or by passing a control signal to a control system for the agent.

[0095] As previously mentioned, components of the system can be trained through reinforcement learning combined with contrastive representation learning. In some implementations, the system maintains a replay buffer to aid in training. The replay buffer stores a number of transitions generated as a result of an agent interacting with an environment. Each transition represents information about the agent's interaction with the environment.

[0096] In these implementations, each transition is an experience tuple that includes: i) a current observation that characterizes the current state of the environment at a point in time; ii) a current action that is performed by the agent in response to the current observation; iii) the next next observation that characterizes the next state of the environment after the agent performs the current action, i.e., the state of the environment transitioned to as a result of the agent performing the current action; and iv) a reward received in response to the agent performing the current action.

[0097] Briefly, in these implementations, RL training can include iteratively sampling one or more batches of transitions from a replay buffer and then training an action selection neural network on the sampled transitions by using a suitable reinforcement learning algorithm. During each RL training iteration, the system can process current observations contained in each sampled transition using an action selection neural network according to current parameter values ​​of the action selection neural network to generate an action selection output, determine a reinforcement learning loss based on the action selection output, and then determine an update to a current value of the action selection network parameter based on computing a gradient of the reinforcement learning loss with respect to the action selection neural network parameter.

[0098] 4 is a diagram illustrating determining updates to parameter values ​​of an attention selection neural network. As shown, at each time step, the system can determine a respective reinforcement learning loss, e.g., RL loss 410A, for the action selection output generated at time step 402A.

[0099] Contrastive representation learning, which can be used to aid RL training to improve training data efficiency, is described further below.

[0100] 3 is a flow chart of an exemplary process 300 for determining updates to parameter values ​​of an attention-selection neural network. For convenience, the process 300 is described as being performed by one or more computer systems located at one or more locations. For example, a suitably programmed reinforcement learning system, such as the reinforcement learning system 100 of FIG. 1, may perform the process 300.

[0101] In particular, the system can iteratively execute process 300 to train the encoder and attention sub-network of the action selection neural network to generate high-quality (e.g., informative and / or predictive) encoded representations and attention sub-network outputs, respectively, that facilitate the generation of high-quality action selection outputs that will result in effective control of the agent in performing a given task.

[0102] The system may perform one iteration of process 300 for each batch of one or more transitions sampled from the replay buffer. At the beginning of each iteration, the system may generate an encoded representation of the current observation contained in each sampled transition by processing the current observation using the encoder sub-network according to the current parameter values ​​of the encoder sub-network. However, unlike at inference time, where the encoded representation is directly provided as input to the attention sub-network, the system generates a masked encoded representation from the encoded representation and then provides the masked encoded representation as input to the attention sub-network.

[0103] As previously mentioned, the encoded representation of the current observation may be in the form of an input vector having a respective input value at each input position of a plurality of input positions in the input sequence. In contrast, the masked encoded representation masks the respective input value at each input position of one or more input positions of the plurality of input positions in the input sequence, i.e., includes a fixed value (e.g., negative infinity, positive infinity, or another predetermined mask value) instead of the original input value at each input position of the one or more input positions.

[0104] To generate a masked encoded representation (hereinafter referred to as a "masked input vector"), the system selects from the encoded representation one or more of the multiple input positions in the input sequence and applies a mask to each input value at each input position of the selected one or more of the multiple input positions in the input sequence, i.e., replaces each input value at each input position of the selected input positions with a fixed value. For example, the selection may be performed via random sampling, and a fixed amount (e.g., 10%, 15%, or 20%) of the input values ​​may be masked for each encoded representation.

[0105] The system processes the masked input vector, which masks the respective input value at each of the one or more input locations in the input sequence, using the attention sub-network according to the current parameter values ​​of the attention sub-network to generate a prediction of the respective input value at each of the one or more input locations in the input sequence (step 302). That is, during contrast learning training, the attention sub-network is trained to perform the auxiliary task of reconstructing the input vector from its masked version.

[0106] The system evaluates (step 304) a contrastive learning objective function that measures the contrastive learning loss (e.g., contrastive loss 420 in FIG. 4) of the attention sub-network in predicting the masked input value from processing the masked input vector.

[0107] Specifically, for each input position of one or more input positions of the multiple input positions in the input sequence, the contrastive learning objective function may measure a first difference between i) a prediction of the respective input value and ii) a respective input value in the input vector corresponding to the encoded representation of the current observation. The first difference may be referred to as a difference evaluated with respect to a "positive example". As illustrated in the example of FIG. 4, at a given time step 402A, the system may determine, for each input position, a respective difference between i) an attention sub-network training output ("X1") 414A including a prediction of the respective input value at that input position and ii) a masked input vector that masks the respective input value originally included in the encoded representation ("Y1") 412A of the observation corresponding to that given time step 402A.

[0108] The control learning objective function may also measure a second difference between i) the prediction of each input value and ii) the respective input value in another input vector corresponding to the encoded representation of the augmented current observation. The second difference may be referred to as a difference evaluated with respect to a "negative example". Additionally or alternatively, the second difference may be a difference between i) the prediction of each input value and ii) the prediction of each input value in another input vector generated by the attention sub-network from the masked input vector corresponding to the augmented current observation. That is, the second difference may be a difference evaluated with respect to the attention sub-network training output generated with respect to the augmented current observation. For example, the first difference and the second difference may be evaluated in terms of Kullback-Leibler divergence.

[0109] In particular, contrastive representation learning typically employs data augmentation techniques to create groupings of data that can be compared to generate meaningful training signals.

[0110] In some implementations, the system may rely on the sequential nature of the input data, and the augmented current observations may be future observations that characterize a future state of the environment after the current state. Additionally, or alternatively, the augmented current observations may be historical observations that characterize past states of the environment prior to the current state.

[0111] As illustrated in the examples of Fig. 4, at a given time step 402A, the system can determine, for each input position, a respective difference between i) an attention sub-network training output ("X1") 414A including a prediction of the respective input value at that input position, and ii) a respective input value in another input vector corresponding to an encoded representation 412B of a future observation received at a future time step 402B. Additionally or alternatively, the system can determine, for each input position, a respective difference between i) an attention sub-network training output ("X1") 414A including a prediction of the respective input value at that input position, and ii) an attention sub-network training output ("X2") 414B including a prediction of the respective input value in the other input vector generated by the attention sub-network from the masked input vector corresponding to the future time step 402B. Specifically, in these examples, each input value in the other input vector can have the same input position as the respective input value in the input vector corresponding to the sampled transition in the other input vector.

[0112] In some other implementations, the system may instead rely on visual representation-based augmentation techniques, where the augmented current observation may be, for example, a geometrically transformed or color space transformed representation of the current observation.

[0113] The system determines updates to the current parameter values ​​of the attention sub-network based on computing gradients of the contrastive learning loss with respect to the attention sub-network parameters (step 306). Additionally, the system determines updates to the current parameter values ​​of the encoder sub-network via backpropagation.

[0114] In some implementations, the system then proceeds to update the current parameter values ​​based on the gradients of the control learning losses by using a conventional optimization algorithm, e.g., an Adam optimization algorithm, including stochastic gradient descent, RMSprop, or the Adam with weight decay ("AdamW") optimization algorithm. Alternatively, the system only proceeds to update the current parameter values ​​once steps 302-306 have been performed for the entire batch of sampled transitions. In other words, the system proceeds to combine the respective gradients determined during a fixed number of iterations of steps 302-306 by computing their weighted or unweighted average, and update the current parameter values ​​based on the combined gradients.

[0115] The system may repeatedly perform steps 302-306 until a control learning training termination criterion is met, e.g., after steps 302-306 have been performed a predetermined number of times or after the gradient of the control learning objective function has converged to a specified value.

[0116] In some implementations, the system can jointly optimize the reinforcement learning loss together with the control learning loss. Thus, in these implementations, the system combines the reinforcement learning loss and the control learning loss, for example, by computing a weighted sum thereof, and then proceeds to update the current parameter values ​​based on the combined loss. In these implementations, steps 302-306 can be performed iteratively until RL training of the system is complete, for example, after the gradient of the reinforcement learning objective function converges to a specified value.

[0117] Figure 5 shows a quantitative example of the performance gains that can be realized by using the control neural network system described herein. Specifically, Figure 5 shows a list of scores (where higher scores indicate greater rewards) received by agents controlled using the control neural network system 110 of Figure 1 for a range of DeepMind Lab tasks. As a platform designed for research and development of general artificial intelligence and machine learning systems, DeepMind Lab (https: / / arxiv.org / abs / 1612.03801) can be used to study how autonomous artificial agents can learn complex tasks in large, partially observed, and visually diverse environments. As shown, it can be seen that the "coberl" agent (corresponding to an agent controlled using the control neural network system described in this specification) generally outperforms the "gtrxl" agent (corresponding to an agent controlled using an existing control system that uses attention mechanisms only, i.e., the "Gated Transformer XL" system described in Parisotto et al., Stabilizing transformers for reinforcement learning, arXiv:1910.06764) by a significant margin for the majority of tasks.

[0118] This specification uses the term "configured" in connection with system components and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that, when activated, causes the system to perform that operation or action. A computer program or programs configured to perform a particular operation or action means that the program or programs contain instructions that, when executed by a data processing device, cause the device to perform that operation or action.

[0119] The subject matter and functional operations described herein may be implemented in digital electronic circuitry, in tangibly embodied computer software or firmware, in computer hardware, including the structures disclosed herein and their structural equivalents, or a combination of one or more of these. The subject matter described herein may be implemented as one or more computer programs, i.e., as one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium for execution by or control of the operation of a data processing device. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random access memory device or a serial access memory device, or a combination of one or more of these storage media. Alternatively, or in addition, the program instructions may be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, generated to encode information for transmission to an appropriate receiving device for execution by the data processing device.

[0120] 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 multiple computers. The apparatus may also be or further include dedicated logic circuitry, such as an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus may optionally include, in addition to the hardware, code that creates an execution environment for a computer program, such as code that constitutes a processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of these.

[0121] A computer program, which may be called or described as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted, or declarative or procedural, and can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may or may 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 other data, e.g., in one or more scripts stored in a markup language document, in a single file dedicated to the program, or in multiple cooperating files, e.g., files that store one or more modules, subprograms, or portions of code. A computer program can be deployed to run on one computer or on multiple computers located at one site or distributed across multiple sites and connected together by a data communication network.

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

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

[0124] The processes and logic flows described herein may be performed by one or more programmable computers executing one or more computer programs to perform functions by manipulating input data and generating output. The processes and logic flows may also be performed by special purpose logic circuitry, e.g., FPGAs or ASICs, or a combination of special purpose logic circuitry and one or more programmed computers.

[0125] A computer suitable for executing a computer program can be based on a general-purpose or dedicated microprocessor, or both, or any other type of central processing unit. In general, the central processing unit receives 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 executing or executing instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory can be supplemented by or incorporated in special-purpose logic circuitry. In general, a computer also includes one or more mass storage devices, e.g., magnetic disks, magneto-optical disks, or optical disks, for storing data, or is operatively coupled to receive data from such mass storage devices, transfer data to such mass storage devices, or both. However, a computer need not have such devices. Additionally, a computer may be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, such as a Universal Serial Bus (USB) flash drive, to name just a few.

[0126] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, non-volatile media, and non-volatile 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 disks and DVD-ROM disks.

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

[0128] A data processing apparatus for implementing machine learning models may also include, for example, dedicated hardware accelerator units for handling the typical, computationally intensive part of machine learning training or production, i.e., the inference workload.

[0129] The machine learning model can be implemented and deployed using a machine learning framework, for example, the TensorFlow framework, the Microsoft Cognitive Toolkit framework, the Apache Singa framework, or the Apache MXNet framework.

[0130] An embodiment of the subject matter described herein can be implemented in a computing system including a back-end component, e.g., as a data server, or in a computing system including a middleware component, e.g., as an application server, or in a computing system including a front-end component, e.g., as a client computer having a graphical user interface, web browser, or app through which a user can interact with an implementation of the subject matter described herein, or in a computing system including a combination of one or more such back-end, middleware, or front-end components. The components of the system can be connected to each other by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0131] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communications network. The relationship between clients and servers arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server sends data, e.g., HTML pages, to a user device, e.g., for the purpose of displaying the data to a user interacting with the device acting as a client, and for the purpose of receiving user input from such a user. Data generated at the user device, e.g., results of user interaction, may be received from the device at the server.

[0132] Although the specification includes many specific implementation details, these should not be construed as limitations on the scope of the invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to certain embodiments of a particular invention. Also, some features described in the context of separate embodiments may also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, although features may be described above as acting in a certain combination, and may even be initially claimed as such, one or more features from a claimed combination may in some cases be removed from the combination, and the claimed combination may be directed to a subcombination or a variation of the subcombination.

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

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

[0135] 100 Reinforcement Learning Systems 102 Agent 104 Environment 106 Action 108 Observations 110 Control Neural Network System 118 Network parameters 120 Action Selection Neural Network 122 Action Selection Output 124 Encoder Sub-Network 126 Coded Representation 128 Attention Subnetwork 130 Attention Sub-Network Output 132 Gating Subnetwork 134 Gating Sub-Network Output 136 Recurrent Subnetworks 140 Action Selection Subnetwork 150 Training Engine 402A, 402B Time Step 412A, 412B Coded representation of observations 414A, 414B Attention subnetwork training output 410 RL loss 420 Control Loss

Claims

1. 1. A system for controlling an agent that interacts with an environment to perform a task, comprising: One or more computers; one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to implement an action selection neural network configured to generate an action selection output that is used to select an action to be performed by the agent; and Equipped with The action selection neural network comprises: an encoder sub-network configured to receive, at each time step of a plurality of time steps, an encoder sub-network input including a current observation characterizing a current state of the environment, and to generate an encoded representation of the current observation; an attention sub-network configured to receive, at each time step of the plurality of time steps, an attention sub-network input including the encoded representation of the current observation, and to generate, at least in part, an attention sub-network output by applying an attention mechanism to the encoded representation of the current observation and to encoded representations of one or more previous observations characterizing one or more previous states of the environment; a recurrent subnetwork configured to, at each time step of the plurality of time steps, receive a recurrent subnetwork input derived from the attention subnetwork output to update a current hidden state of the recurrent subnetwork corresponding to the time step, and to generate a recurrent subnetwork output, the current hidden state being generated by processing a previous recurrent subnetwork input, and the update of the current hidden state modifying the current hidden state by processing the received recurrent subnetwork input; an action selection subnetwork configured to receive, at each time step of the plurality of time steps, an action selection subnetwork input including the recurrent subnetwork output, and to generate an action selection output used to select the action to be performed by the agent in response to the current observation; A system comprising:

2. The system of claim 1 , wherein the encoded representation of the current observation comprises an input vector having a respective input value at each input position of a plurality of input positions in an input order.

3. The attention sub-network comprises: receiving an attention layer input for each layer input location of the plurality of layer input locations; and For each particular layer input position in the layer input order, applying the attention mechanism to the attention layer inputs at the particular layer input positions using one or more queries derived from the attention layer inputs at the particular layer input positions to generate respective attention layer outputs for the particular layer input positions.

3. The system of claim 2, comprising a plurality of attention layers, each configured to:

4. The system of claim 3 , wherein the attention mechanism is a masked attention mechanism.

5. 5. The system of claim 1, wherein the recurrent sub-network comprises one or more long short-term memory (LSTM) layers.

6. 6. The system of claim 1, wherein the action selection output includes a Q-value for each action in a possible set of actions that is an estimate of a return that would be received if the agent performed the action in response to the current observation.

7. 7. The system of claim 1, wherein the action selection neural network further comprises a gating layer configured to apply a gating mechanism to i) the encoded representation of the current observation and ii) the attention sub-network output to generate the recurrent sub-network input.

8. 8. The system of claim 7, wherein applying the gating mechanism to i) the encoded representation of the current observation and ii) the attention sub-network output comprises applying a gated recurrent unit (GRU) to i) the encoded representation of the current observation and ii) the attention sub-network output.

9. 9. The system of claim 1, wherein at each time step of the plurality of time steps, the attention sub-network input comprises the encoded representation of the current observation and the encoded representation of the one or more previous observations characterizing the one or more previous states of the environment.

10. 10. One or more computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to implement the action selection neural network of any one of claims 1 to 9.

11. 10. A method comprising operations that the action selection neural network of any one of claims 1 to 9 is configured to perform.

12. processing a masked input vector that masks the respective input value at each of the one or more input positions of a plurality of input positions in an input sequence to generate a prediction of the respective input value at each of the one or more input positions of the plurality of input positions in an input sequence using at least the attention sub-network having a plurality of attention sub-network parameters; for each input position of the one or more input positions of the plurality of input positions in the input sequence, a first difference between i) the prediction of the respective input value and ii) the respective input value in the input vector included in the encoded representation of the current observation; and a second difference between i) the prediction of the respective input value and ii) the respective input value in an input vector included in the encoded representation of the augmented current observation; Evaluating a contrastive learning objective function that measures determining updates to current values ​​of the plurality of attention sub-network parameters based on the computed gradient of the contrast learning objective function; 10. A method for training the action selection neural network of claim 1, comprising:

13. The masked input vector is randomly selecting the one or more input positions of the plurality of input positions in the input sequence; and applying a mask to the respective input value at each of the randomly selected one or more of the plurality of input positions in the input sequence. The method of claim 12 , further comprising the step of generating by:

14. 14. The method of claim 12, wherein the augmented present observations include future observations characterizing a future state of the environment subsequent to the current state.

15. 14. The method of claim 12, wherein the augmented current view comprises a geometrically transformed current view or a color space transformed current view.

16. processing the current observation using the action selection neural network having a plurality of action selection network parameters to generate the action selection output; determining a reinforcement learning loss based on the action selection output; determining updates to current values ​​of the action selection network parameters based on the reinforcement learning losses; The method of any one of claims 11 to 15, further comprising:

17. 1. A computer-implemented method for controlling an agent that interacts with an environment to perform a task, comprising: At each time step of the multiple time steps, receiving an encoder sub-network input comprising a current observation characterizing a current state of the environment; generating an encoded representation of the current observation; generating an attention sub-network output, at least in part, by applying an attention mechanism to the encoded representation of the current observation and to encoded representations of one or more previous observations characterizing one or more previous states of the environment; updating a current hidden state of the recurrent subnetwork corresponding to the time step based on a recurrent subnetwork input derived from the attention subnetwork output to generate a recurrent subnetwork output, the current hidden state being generated by processing a previous recurrent subnetwork input, and the current hidden state update modifying the current hidden state by processing the recurrent subnetwork input; generating an action selection output based on action selection sub-network inputs including the recurrent sub-network output; selecting an action to be performed by the agent based on the action selection output; sending control data to said agent instructing said agent to perform said selected action; A computer-implemented method comprising:

18. One or more computers; one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform the respective operations of any one of the methods according to any one of claims 12 to 17; A system comprising:

19. One or more computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform each of the operations of any one of the methods recited in any one of claims 12 to 17.

Images