Autoregressive generation of a sequence of data elements that defines the actions to be performed by the agent

Abstract

A method, system, and apparatus, including a computer program encoded on a computer storage medium, for selecting an action to be performed by an agent to interact with an environment using an action selection neural network. In one aspect, the method includes generating, at each time step in a sequence of time steps, a current representation of a state of a task being performed by the agent in the environment as of the current time step as a sequence of data elements, autoregressively generating a sequence of data elements representing a current action to be performed by the agent at the current time step, and having the agent perform the current action at the current time step after autoregressively generating the sequence of data elements representing the current action.

Description

[Technical Field]

[0001] This specification relates to processing data using machine learning models. [Background technology]

[0002] A machine learning model receives input and generates an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models, which generate an output based on the received input and the values of the parameters of the model.

[0003] Some machine learning models are deep models that employ multiple layers of the model to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a nonlinear transformation to a received input to generate an output. [Prior art documents] [Non-patent literature]

[0004] [Non-Patent Document 1] K. He et al., “Identity mappings in deep residual networks,” European Conference on Computer Vision, pp. 630–645, 2016. [Non-patent document 2] A. Dosovitskiy et al., “An image is worth 16x16 words: transformers for image recognition at scale,” arXiv:2010.11929v2, 2021 [Non-patent document 3] A. Vaswani et al., "Attention is all you need," arXiv:1706.03762v5, December 6, 2017 [Non-patent document 4] "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context," Dai et al., arXiv:1901.02860v3, June 2, 2019 Summary of the Invention [Means for solving the problem]

[0005] This specification describes an action selection system, implemented as a computer program on one or more computers at one or more locations, for controlling agents that interact with an environment to perform tasks.

[0006] Throughout this specification, a "data element" may refer, for example, to a numeric value (e.g., an integer or floating-point value) or to an embedding. An embedding refers to an ordered collection of numeric values, for example, a vector, matrix, or other tensor of numeric values.

[0007] According to a first aspect, there is provided a method implemented by one or more computers for selecting an action to be performed by an agent to interact with an environment using an action selection neural network, particularly a trained action selection neural network. The method includes, at each time step in a sequence of time steps, generating, e.g., from current observations characterizing the state of the environment, a current representation of the state of a task being performed by the agent in the environment as of the current time step as a (first) sequence of data elements. The method also includes autoregressively generating a (second) sequence of data elements representing a current action to be performed by the agent at the current time step. For example, the (second) sequence of data elements may comprise a plurality of action data elements that collectively represent an action to be performed by the agent. In an implementation, the step of autoregressively generating the (second) sequence of data elements includes: for each position (in the second sequence of data elements), starting from the first position in the sequence of data elements representing the current action, processing a current representation of the state of the task using an action selection neural network to generate a score distribution over the set of possible data elements; selecting a data element for the position in the sequence of data elements representing the current action according to the score distribution; and updating the current representation of the state of the task by concatenating the selected (action) data element for the position with the current representation of the state of the task. That is, the updated current representation of the state of the task, i.e., the (first) sequence of data elements, is updated for the autoregressive generation of the (second) sequence of data elements, specifically, by processing the current (currently updated) representation of the state of the task to select an (action) data element for the next position. After autoregressively generating the sequence of data elements representing the current action, the method causes the agent to execute the current action at the current time step.The method may then update the current representation of the state of the task using the current observations for the next time step.

[0008] In some implementations, for each time step in the sequence of time steps, generating a current representation of the state of the task as of the current time step includes receiving current observations characterizing the state of the environment at the current time step, generating a representation of the current observations as a sequence of data elements, and including the representation of the current observations as a sequence of data elements in the current representation of the state of the task as of the current time step, for example, by concatenating a (first) sequence of data elements representing the current state of the task and the representation of the current observations as a sequence of data elements.

[0009] In some implementations, the current observation is defined by a collection of numerical values, and generating a representation of the current observation as a sequence of data elements comprises concatenating each numerical value in the collection of numerical values that defines the current observation into a sequence of numerical values in a predefined order, i.e., that defines an ordering of the numerical values of the observation.

[0010] In some implementations, generating a representation of the current observation as a sequence of data elements further includes discretizing each number in the collection of numbers that defines the current observation.

[0011] In some implementations, the current observation characterizing the current state of the environment at the current time step comprises an image of the environment defined by an array of pixels.

[0012] In some implementations, generating a representation of the current observation as a sequence of data elements includes combining a target return to be achieved by the agent's interaction with the environment with the representation of the current observation as a sequence of data elements, the target return defining a cumulative measure of reward to be achieved as a result of the agent's interaction with the environment.

[0013] In some implementations, for each time step after the first time step in the sequence of time steps, including a representation of the current observation as a sequence of data elements in a current representation of the state of the task as of the current time step includes receiving a representation of the state of the task as of the previous time step as a sequence of data elements, and concatenating the representation of the current observation as a sequence of data elements with the representation of the state of the task as of the previous time step as a sequence of data elements to generate a current representation of the state of the task as of the current time step.

[0014] In some implementations, the representation of the state of the task at the previous time step represents, for each time step preceding the current time step, (i) each observation that characterizes the state of the environment at the time step, and (ii) each action performed by the agent at the time step.

[0015] In some implementations, at a first time step in the sequence of time steps, including a representation of the current observation as a sequence of data elements in a current representation of the state of the task as of the current time step includes receiving a prompt comprising data characterizing a task to be performed by the agent in the environment; generating a representation of the prompt as a sequence of data elements; and concatenating the representation of the current observation as a sequence of data elements with the representation of the prompt as a sequence of data elements to generate a current representation of the state of the task as of the current time step.

[0016] In some implementations, the prompt comprises one or more of a demonstration of the task, a goal observation that characterizes a goal state of the environment, or a sequence of text in natural language that provides instructions related to the task.

[0017] In some implementations, the method includes training the action selection neural network such that the action selection neural network is trained on a set of training examples, where for each training example, the training example is represented as a sequence of data elements, at least one of the data elements in the sequence of data elements representing the training example is designated as an action data element, and training the action selection neural network on the training examples generates the action data element included in the training example.

[0018] In some implementations, the set of training examples includes respective training examples from multiple different control domains, each control domain being associated with (i) a corresponding agent, (ii) a corresponding environment, and (iii) a corresponding task, and each training example from each control domain characterizes an interaction between the corresponding agent and the corresponding environment by performing an action to accomplish the corresponding task.

[0019] In some implementations, the plurality of different control domains includes a first control domain in which corresponding environmental observations have a first dimensionality and a second control domain in which corresponding environmental observations have a second, different dimensionality.

[0020] In some implementations, the plurality of different control domains includes a first control domain in which actions performed by corresponding agents have a first dimensionality, and a second control domain in which actions performed by corresponding agents have a second, different dimensionality.

[0021] In some implementations, the set of training examples includes multiple language modeling training examples, each language modeling training example representing a sequence of text in a natural language.

[0022] In some implementations, the action selection neural network comprises multiple self-attention neural network layers. Generally, a self-attention neural network layer has an attention layer input for each element of the input and is configured to apply an attention mechanism across the attention layer inputs to generate an attention layer output for each element of the input. There are many different attention mechanisms that can be used.

[0023] In some implementations, starting from the first position in the sequence of data elements representing the current action, for each position, selecting a data element for the position includes selecting the data element having the highest score under the score distribution.

[0024] In some implementations, for each time step in the sequence of time steps, the sequence of data elements representing the state of the task at the time of the current time step comprises a sequence of numerical values, a sequence of embeddings, or a sequence containing numerical values in some positions and embeddings in other positions.

[0025] In some implementations, the agent is a mechanical agent that interacts with a real-world environment. Thus, the selected action may be an action performed by the mechanical agent in the real-world environment, such as an action that causes the mechanical agent to physically manipulate one or more objects in the environment, and the observations characterizing the state of the environment may be observations of the real-world environment. The observations may be multimodal observations. The method may use an action selection neural network to perform one or more tasks, and a particular advantage of the described system is that the same action selection neural network with the same set of parameters (weights) can be used to perform multiple different tasks. In some implementations, the system, and in particular the action selection neural network, has over 1.2 billion learnable parameters, enabling the ability to perform multiple different tasks.

[0026] In some implementations, for each time step in the sequence of time steps, the sequence of data elements representing the state of the task at the time of the current time step comprises a sequence of numerical values, a sequence of embeddings, or a sequence containing numerical values in some positions and embeddings in other positions.

[0027] In some implementations, the current observation comprises an image, and generating a representation of the current observation as a sequence of data elements includes generating a respective initial patch embedding corresponding to each of a plurality of patches in the image, and processing the initial patch embeddings using an encoder neural network to generate a respective final patch embedding for each of the plurality of patches in the image, each final patch embedding being included as a respective data element in the sequence of data elements representing the current observation.

[0028] In some implementations, generating each initial patch embedding corresponding to a patch in the image includes generating pixel embeddings that represent pixels in the patch in the image, generating patch position embeddings that represent the position of the patch in the image, and generating the initial patch embedding for the patch by combining the pixel embedding and the patch position embedding for the patch.

[0029] In some implementations, the encoder neural network comprises one or more self-attention neural network layers.

[0030] In some implementations, the encoder neural network comprises one or more residual blocks.

[0031] In some implementations, the agent is a mechanical agent that interacts with a real-world environment.

[0032] In some implementations, selecting an action to be performed by the mechanical agent includes selecting an action to cause the mechanical agent to physically manipulate one or more objects in the environment.

[0033] According to another aspect, a system is provided, the system comprising one or more computers and one or more storage devices communicatively coupled to the one or more computers, the one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the methods described herein.

[0034] One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations of the methods described herein.

[0035] Particular embodiments of the subject matter described in this specification can be implemented to realize one or more of the following advantages.

[0036] The action selection system described herein uses an autoregressive action selection neural network that operates on sequences of data elements to select actions to be performed by an agent in an environment. Specifically, the action selection system represents both observations and actions as sequences of data elements and uses an action selection neural network to operate on these sequences to autoregressively generate sequences of data elements that represent actions to be performed by the agent in the environment. Because the action selection neural network operates on sequences of data elements, it can be trained on any training examples that can be represented as sequences of data elements. Thus, the action selection neural network can be trained on training examples that represent the interaction of any agent with any environment to perform any task, regardless of the dimensionality of the observations of the environment and the actions performed by the agent, respectively.

[0037] The action selection system trains the action selection neural network on a highly diverse set of training examples representing interactions between multiple different agents and multiple different environments to perform multiple different tasks. Thus, the action selection neural network learns a flexible and transferable understanding of agent control that allows the action selection neural network to generalize quickly and effectively to new domains. In particular, the action selection neural network can perform "few-shot learning," i.e., the action selection neural network can be trained to achieve an acceptable level of performance on a task in a new domain after being trained on only a few training examples from the new domain. In some cases, the action selection neural network can perform "zero-shot learning," i.e., by achieving an acceptable level of performance on a task in a new domain without being trained on any training examples from the new domain. Thus, the action selection system provides a generic model for agent control that is more broadly applicable than conventional action selection systems. The action selection system requires less training data and fewer training iterations than conventional systems to achieve an acceptable level of performance in controlling agents in new domains, thereby enabling more efficient use of computational resources (e.g., memory and computing power).

[0038] In addition to training the action selection neural network to perform agent control tasks, the action selection system can also train the action selection neural network to perform language modeling, i.e., by training the action selection neural network on sequences of data elements representing text in natural language. Training the action selection neural network to perform language modeling can accelerate training and improve performance of the action selection neural network, for example, by improving the ability of the action selection neural network to implicitly infer the meaning of natural language prompts provided to the action selection neural network.

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

[0040] [Figure 1] FIG. 1 illustrates an exemplary action selection system. [Figure 2] FIG. 1 illustrates training examples from different domains. [Figure 3A] FIG. 1 illustrates operations performed by an action selection system to select actions to be performed by an agent interacting with an environment to accomplish a task. [Figure 3B] FIG. 1 illustrates operations performed by an action selection system to select actions to be performed by an agent interacting with an environment to accomplish a task. [Figure 4] FIG. 1 is a flow diagram of an example process for selecting an action to be performed by an agent to interact with an environment in a current time step. DETAILED DESCRIPTION OF THE INVENTION

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

[0042] 1 illustrates an exemplary action selection system 100. Action selection 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.

[0043] The system 100 selects an action 102 to be performed by an agent 120 interacting with the environment 118 at each of a number of time steps to accomplish a task within the environment 118 .

[0044] At each time step, the system 100 receives observations 116 that characterize the current state of the environment 118 and, in response to the observations 116, selects actions 102 to be performed by the agent 120. As will be explained later, the actions 102 at a time step may be represented by a sequence of action data elements.

[0045] Each time step may be associated with a reward based, for example, on the state of the environment 118 at the time step, the action 102 performed by the agent 120 at the time step, or both. In general, the reward may be expressed as a numerical value. The reward may be based on any event in the environment 118 or aspects of the environment 118. For example, the reward may indicate whether the agent 120 accomplished a task in the environment (e.g., navigating to a goal location in the environment 118) or the agent's progress toward completing the task. In some implementations, the reward may be a sparse reward, having a value of 0 at each time step before the agent accomplishes the task and a value of 1 (or some other positive value) at the time step in which the agent accomplishes the task. In some implementations, the reward may be a dense reward, having a non-zero value at time steps before the agent accomplishes the task; for example, if the task involves navigating to a goal location, the reward at each time step may vary continuously based on the agent's proximity to the goal location.

[0046] The training engine 112 can train the system 100 to select actions that, by performing the actions 102 selected by the system 100, increase the "profit" generated by the interaction of the agent 120 with the environment 118, as described in more detail below. Profit refers to a cumulative measure of rewards, e.g., a time-discounted sum of rewards, generated by the interaction of the agent 120 with the environment 118.

[0047] In some implementations, the environment is a real-world environment and the agent is a mechanical agent that interacts with the real-world environment. For example, the agent may be a robot that interacts with the environment to perform a task, such as to place an object of interest in the environment, move the object of interest to a specified location in the environment, physically manipulate the object of interest in a specified manner in the environment, or navigate to a specified destination in the environment, or the agent may be an autonomous or semi-autonomous land vehicle, aircraft, or watercraft that navigates through the environment to a specified destination in the environment. In a particular example, the agent may be a robot that interacts with the real-world environment using a mechanical gripping tool, for example, to stack a set of objects (e.g., boxes) in the environment or to assemble a set of components (e.g., electronic components) in the environment.

[0048] In these implementations, the observations may include, for example, one or more of images (where the images may be represented, for example, as an array of pixels), object position data, and sensor data, for example, from image sensors, distance sensors, or position sensors, or from actuators, to capture the observations as the agent interacts with the environment.

[0049] For example, in the case of a robot, the observations may include data characterizing the current state of the robot, such as one or more of joint positions, joint velocities, joint forces, torques or accelerations, e.g., gravity-compensated torque feedback, and the global or relative pose of an item held by the robot.

[0050] 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.

[0051] The observations may also include, for example, data acquired by one or more sensor devices that sense the real-world environment, e.g., sensed electronic signals such as motor current or temperature signals, and / or image or video data, e.g., from a camera or LIDAR sensor, e.g., data from sensors on the agent or from sensors located remotely from the agent in the environment.

[0052] In the case of electronic agents, the observations may include data from one or more sensors monitoring portions of the 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.

[0053] An action may be a control input for controlling a robot, e.g., torque for a joint of the robot, or a higher level control command, or a control input for controlling an autonomous or semi-autonomous land vehicle or aircraft or watercraft, e.g., torque to a control surface or other control element of the vehicle, or a higher level control command.

[0054] In other words, an action may include, for example, position, velocity, or force / torque / acceleration data for one or more joints of a robot or part of another mechanical agent. An action may additionally 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, the control of which affects the observed state of the environment. For example, in the case of an autonomous or semi-autonomous land vehicle, aircraft, or watercraft, an action may include actions for controlling navigation, e.g., steering, and movement, e.g., braking and / or accelerating the vehicle. As explained above, an action at a particular time step may have multiple components, each represented by a respective action data element.

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

[0056] For example, 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 navigates through the motion simulation. In these implementations, the actions may be control inputs for controlling a simulated user or a simulated vehicle.

[0057] In another example, the simulated environment may be a video game and the agent may be a simulated user playing the video game.

[0058] In a further example, the simulated environment may be a protein folding environment, where each state is a respective state of a protein chain and the agent is a computer system for determining how to fold the protein chain. In this example, the actions are possible folding actions for folding the protein chain, and the tasks to be performed may include, for example, folding the protein so that it is stable and so that it achieves a particular biological function.

[0059] In general, for a simulated environment, the observations may include simulated versions of one or more of the previously described observations or types of observations, and the actions may include simulated versions of one or more of the previously described actions or types of actions.

[0060] In some cases, the action selection system 100 may be used to control an agent's interaction with a simulated environment, and the training engine 112 may train parameters of the action selection system based on the agent's interaction with the simulated environment. After the action selection system is trained based on the agent's interaction with the simulated environment, the agent may be deployed in a real-world environment, and the trained action selection system may be used to control the agent's interaction with the real-world environment. Training the action selection system based on the agent's interaction with a simulated environment (i.e., instead of the real-world environment) can avoid agent attrition and reduce the possibility that the agent may damage itself or aspects of its environment by performing poorly chosen actions.

[0061] In some other applications, agents may control actions in a real-world environment containing items of equipment, for example, in a data center or grid mains power or water distribution system, or in a manufacturing plant or service facility. Observations may then relate to the operation of the plant or facility. For example, observations may include observations of power or water use by equipment, or observations of power generation or distribution control, or observations of resource use or waste generation. The agents may perform tasks to increase efficiency, for example, by reducing resource usage, and / or control actions in the environment to reduce the environmental impact of operations in the environment, for example, by reducing waste. Actions may include actions that control or impose operating conditions on items of equipment in the plant / facility, for example, to adjust or turn on / off components of the plant / facility, and / or actions that result in changes to settings in the operation of the plant / facility.

[0062] In some further applications, the environment is a real-world environment and the agent manages the distribution of tasks across computing resources, for example, on mobile devices and / or in a data center. In these implementations, the actions may include assigning tasks to specific computing resources, and the tasks to be performed may include minimizing the time required to complete the set of tasks using the designated computing resources.

[0063] As a further example, the action may include presenting an advertisement, the observation may include an impression or click-through count or rate of the advertisement, and the reward may characterize previous selections of items or content taken by one or more users. In this example, the task to be performed may include maximizing selections of items or content by one or more users.

[0064] As a further example, an agent can generate actions that represent sequences of text in a natural language. In some implementations, a task can be, for example, to generate a sequence of natural language text responsive to an observation represented by the sequence of natural language text. In some implementations, a task can be, for example, to generate a sequence of natural language text (e.g., "turn left," "speed up," "turn on the lights," etc.) that represents instructions for controlling a (real or simulated) physical agent to perform a task in a (real or simulated) physical environment.

[0065] As a further example, an agent may generate actions that represent sequences of computer code (e.g., in a computer programming language). In some implementations, a task may involve receiving observations that define a natural language description of desired computer code and, in response, generating a sequence of computer code that conforms to the natural language description of the desired computer code. In some implementations, a task may involve receiving observations that define an input sequence of computer code and, in response, generating an output sequence of computer code that is a completion of the input sequence of computer code (e.g., a logical extension of the input sequence of computer code).

[0066] To select actions 102 to be performed by agent 120, system 100 maintains and iteratively updates a current task state 110, which is represented as a sequence of data elements. The sequence of data elements representing the current task state 110 may be, for example, a sequence of numbers, a sequence of embeddings, or a sequence that includes numbers in some positions and embeddings in other positions. At each time step, the current task state 110 represents the state of the task being performed by the agent in the environment at the time of that time step.

[0067] In some cases, prior to the first time step (i.e., the first time step in a sequence of multiple time steps during which the agent interacts with the environment to perform a task), system 100 may initialize current task state 110 with a “prompt,” which may be any suitable data that characterizes the task to be performed by agent 120 within environment 118. The prompt may be provided to system 100, for example, by a user of system 100. Some examples of prompts are described in more detail next, and exemplary techniques for representing a prompt as a sequence of data elements are described in more detail below.

[0068] In some implementations, a prompt may include a demonstration of a task to be performed by an agent within the environment. That is, a prompt may characterize the agent's interaction with the environment over a sequence of time steps during which the agent progresses toward accomplishing a task within the environment. A prompt may be defined by a sequence of "dialogue tuples," where each dialogue tuple corresponds to a respective time step and represents an observation of the environment at the time step, an action performed by the agent at the time step, or both.

[0069] The prompt may include a demonstration of a task that is different from (but related to) the task to be performed by agent 120 in environment 118. For example, if agent 120 is a robotic agent and the task to be performed by agent 120 involves grasping and moving a certain type of object (e.g., an object having a cuboid shape), the prompt may define a demonstration of grasping and moving a different type of object (e.g., an object having a spherical shape).

[0070] The prompt may include a demonstration of the task in an environment different from the environment 118 in which the agent 120 will perform the task. For example, if the agent 120 is a domestic robotic agent and the task to be performed by the agent involves cleaning a room (e.g., the kitchen), the prompt may define a demonstration of the agent cleaning a different room (e.g., the bathroom).

[0071] The prompt may include a demonstration of a task being performed by an agent different from the agent 120 being controlled by the system 100. For example, if the agent is a robotic agent 120 with a robotic arm, the prompt may define a demonstration by an agent with a differently configured (e.g., having a different length) robotic arm performing the task.

[0072] In some implementations, the prompt may include, for example, a "goal" observation that characterizes a goal state of the environment such that agent 120 accomplishes the task by performing an action that transitions the environment to the goal state (or a state related to the goal state). For example, if agent 120 is a robotic agent and the task to be performed by agent 120 involves assembling a set of components (e.g., electronic or mechanical components), the goal observation may be, for example, an image showing the set of components assembled in a desired configuration.

[0073] In some implementations, the prompt may include a sequence of text in a natural language (e.g., English) that provides instructions related to a task to be performed by the agent 120 within the environment 118. For example, if the agent 120 is a semi-autonomous or fully autonomous vehicle, the prompt may be a sequence of words such as "park the car in the parking lot closest to the store entrance" or "merge onto the highway and move into the left-most lane."

[0074] In some implementations, the prompt may include data that characterizes the task to be performed by agent 120 in a number of different ways; for example, the prompt may include both a demonstration of the task and a sequence of text in natural language that provides instructions related to the task.

[0075] System 100 may represent a prompt as a sequence of data elements in any suitable manner. For example, if the prompt includes a sequence of text, system 100 may represent the sequence of text as a sequence of tokens from a predefined set of tokens and then map each token to a corresponding numeric value according to a predefined mapping. The set of tokens may include, for example, characters, n-grams, word pieces, words, or combinations thereof. Exemplary techniques for representing observations and actions as sequences of data elements that may be applied to represent task demonstrations or objective observations within a prompt as sequences of data elements are described in more detail below.

[0076] In general, a prompt encodes information that can enable system 100 to infer the task to be performed by agent 120 in environment 118 and, therefore, select the actions to be performed by agent 120 to accomplish the task. The system may be able to infer the task to be performed from the format of the observations and actions in the representation of the state of the task. However, the system may need further context to disambiguate the task, which can be provided by a prompt.

[0077] In some cases, as described above, the prompt represents a task demonstration in a different environment (e.g., different from environment 118), or a task demonstration by a different agent (e.g., configured differently from agent 120), or a task demonstration of a task that is different but related to the task to be performed by agent 120. In these cases, system 100 can incorporate information encoded in the prompt along with information encoded in observations 116 of environment 118 to infer the task to be performed by agent 120 in environment 118. For example, if the agent is a domestic robot and the task to be performed by the agent involves cleaning a target house, the prompt may include a demonstration of the agent cleaning a different house. In this example, the system can combine information representing the cleaning task, as encoded in the prompt, with information representing the target house, as encoded in observations received by the system, to implicitly infer that the task to be performed by the agent involves cleaning the target house.

[0078] In some cases, system 100 can implicitly infer the task to be performed by agent 120, even in the absence of a prompt, particularly based on information encoded in observations received by system 100. For example, system 100 can implicitly infer from the observations that agent 120 is interacting with a type of environment in which the agent generally performed a particular task during training, and based on this, select the action to be performed by the agent to accomplish the particular task.

[0079] At each time step, system 100 receives current observations 116 that characterize the state of environment 118 at the time step and uses the current observations 116 to update current task state 110. For example, system 100 may represent current observations 116 as a sequence of data elements and update current task state 110 by concatenating the sequence of data elements representing the current observations with the sequence of data elements representing current task state 110. That is, in this example, the updated task state is represented by a sequence of data elements defined by concatenating (i) the sequence of data elements representing current task state 110 and (ii) the sequence of data elements representing current observations 116.

[0080] The system 100 may represent the current observation 116 for a time step as a sequence of data elements in any suitable manner. The sequence of data elements representing the current observation 116 may be, for example, a sequence of numeric values, a sequence of paddings, or a sequence that includes numeric values in some positions and paddings in other positions. Several exemplary techniques for representing the current observation 116 as a sequence of data elements are now described.

[0081] Generally, when received by system 100, current observation 116 is defined by an ordered collection of numbers, e.g., a vector, matrix, or other tensor of numbers. (The number of numbers in an ordered collection of numbers defining an observation is sometimes referred to as the "dimensionality" of the observation.) In some implementations, system 100 can represent the collection of numbers defining current observation 116 as a sequence of numbers, e.g., by concatenating the numbers defining the current observation into a sequence of numbers in an arbitrary but fixed order.

[0082] For example, if observation 116 includes an image represented by an array of pixel intensity values, system 100 may represent the array of pixel intensity values as a sequence of numerical values by concatenating each pixel intensity value in the array of pixel intensity values representing the image into a sequence of pixel intensity values in an arbitrary but fixed order. If the array of pixel intensity values is an N×N array having N rows and N columns, the system may aggregate the pixel intensity values in each row of the array, for example, from the first row in the array to the last row in the array, starting from the first position in the row to the last position in the row.

[0083] As another example, if an observation includes a position value (e.g., representing the position of an agent within the environment), a velocity value (e.g., representing the velocity of the agent within the environment), and an acceleration value (e.g., representing the acceleration of the agent within the environment), the system may concatenate these values in any predefined order, such as, for example, the position value followed by the velocity value followed by the acceleration value.

[0084] In some implementations, to generate a representation of an observation 116 (or some portion of an observation 116) as a sequence of numerical values, the system 100 first generates an embedding (e.g., a lower-dimensional embedding) of the observation 116 by processing the observation 116 using an encoder machine learning model. The system 100 can then concatenate the numerical values that define the embedding of the observation into a sequence of numerical values that represents the observation, in an arbitrary but fixed order. The encoder machine learning model can be, for example, an encoder neural network of an autoencoder machine learning model.

[0085] In some implementations, the system 100 can generate a representation of an observation as a sequence of one or more embeddings. For example, the system 100 can generate a representation of an image as a sequence of embeddings by dividing an image into a sequence of patches and then using an encoder machine learning model to generate a respective embedding for each patch. The system can then concatenate the respective embeddings of the image patches to generate a representation of the image as a sequence of embeddings. The encoder machine learning model can be implemented as a neural network having any suitable neural network architecture. Some examples of possible architectures for an encoder neural network, i.e., for implementing an encoder machine learning model, are described next.

[0086] In one example, the encoder neural network may have a residual neural network architecture that includes a sequence of residual blocks, e.g., where the input to each residual block is added to the output of the residual block. In a particular example, the encoder neural network may be implemented using a v2 ResNet architecture, as described with reference to, e.g., K. He et al., “Identity mappings in deep residual networks,” European Conference on Computer Vision, pp. 630-645, 2016. The encoder neural network may be configured to receive a respective initial embedding representing each patch in the image. The initial embedding of an image patch may be based on (i) a patch pixel embedding representing the pixels in the image patch, e.g., generated by concatenating the pixels in the image patch into a vector, and (ii) a patch position embedding representing the position of the patch in the image. For example, the initial embedding of an image patch may be the sum or concatenation of the patch pixel embedding and the patch position embedding for the patch. The encoder neural network may be configured to process the initial embedding of each image patch to generate a final embedding for the patch.

[0087] The system can generate a patch position embedding for an image patch, i.e., a representation of the position of the image patch in the image from which it was extracted, in any suitable manner. For example, the system can generate relative row and column intervals for the patch by normalizing the pixel interval of the patch by the image resolution. The system can quantize the normalized row and column intervals into a finite vocabulary of indices that index (i) a table that stores row position encodings and (ii) a table that stores column position encodings. The system can take the indexed row and column position encodings and add (or possibly combine) the row and column position encodings to produce the patch position embedding.

[0088] As another example, the encoder neural network may have an attention-based neural network architecture. For example, the encoder neural network may include one or more self-attention neural network layers. The encoder neural network may, for example, use the self-attention neural network layers to iteratively update the initial embeddings of the image patches (as described above) to generate a respective final embedding for each image patch. In a particular example, the encoder neural network may have a Vision Transformer architecture, such as that described with reference to A. Dosovitskiy et al., "An image is worth 16x16 words: transformers for image recognition at scale," arXiv:2010.11929v2, 2021.

[0089] The encoder neural network may be jointly trained with the action selection neural network 108, for example, to optimize a loss function. For example, the training engine 112 may backpropagate the gradient of the loss function through the action selection neural network 108 and into the encoder neural network.

[0090] In some cases, an observation for a time step may include multiple constituent observations. For example, an observation for a time step may include respective images captured by multiple camera sensors of an agent. In these cases, as part of representing the observation as a sequence of data elements, the system may combine (e.g., add) a respective "observation-level" position embedding with each data element in the sequence. The observation-level position embedding for a data element may, for example, characterize the index of the constituent observation represented by the data element. The observation-level position embedding may, for example, be a learned or pre-defined embedding. An action data element for a time step may be combined with an action embedding, which may be the same for each action data element.

[0091] In some cases, system 100 may perform "return conditioning" by generating an additional numerical value (a "return value") that represents a target return to be achieved by the interaction of agent 120 with environment 118 (at one or more time steps) and combining the return value with a sequence of data elements that represent the current observation. For example, system 100 may perform return conditioning by concatenating the return value with a sequence of data elements that represent the current observation 116.

[0092] In general, the system 100 aims to select an action 102 that maximizes the profit received by the agent 120. Thus, the system 100 can set the profit value to a predefined “expert” profit value that represents the profit that would be achieved by the agent through expert execution of the task. The system can calculate the expert profit value, for example, as the average profit value achieved when the agent performs the task one or more times under the control of an expert, e.g., a human expert. Performing profit adjustments enables the training engine 112 to effectively train the system 100 on training examples in which the agent receives a variety of possible profits (including potentially low profits), as described in more detail below. The training engine 112 can normalize the profit values used during training so that the optimal profit value for each task is the same predefined value, e.g., a value of 1, by dividing the profit value by the expert profit value, as described in more detail below.

[0093] In some cases, when system 100 initializes current task state 110 using a prompt that includes a task demonstration, the prompt may include a revenue adjustment as well. For example, each observation in the prompt may include an additional value that represents the revenue achieved during the task demonstration.

[0094] As part of representing the current observation 116 as a sequence of data elements, the system 100 may discretize each numeric value in the collection of numeric values that defines the current observation 116. Discretizing a numeric value may refer to mapping the numeric value to a corresponding numeric value from a finite, predefined set of “discretized” numeric values, e.g., integer values in the range [0, 255]. To discretize a numeric value, the system 100 may first apply a transformation function (e.g., a μ-law transformation function) to the numeric value to map the numeric value to a predefined range (e.g., the range [−1, 1]). The predefined range may be associated with a predefined partition into a set of intervals, and each of these intervals may be associated with a corresponding discretized numeric value from the predefined set of discretized numeric values. Applying the transformation function to a numeric value will cause the numeric value to fall within one of the intervals, and the system 100 may discretize the numeric value by mapping the numeric value to a discretized numeric value associated with the interval.

[0095] After updating the current task state 110 using the current observations 116, the system 100 processes the current task state 110 to autoregressively generate a sequence of one or more data elements that collectively represent an action 102 to be performed by the agent at the current time step. Each data element in the sequence of data elements representing an action 102 is referred to herein as an "action data element" 104, i.e., the action 102 becomes defined by the sequence of action data elements 104 generated by the action selection neural network 108.

[0096] The system 100 generates, in order, starting from the first position in the sequence, a respective action data element 104 at each position in the sequence of action data elements defining the current action 102. The system 100 generates each action data element 104 by processing the current task state 110 using the action selection neural network 108 according to parameter values of the action selection neural network 108 to generate a score distribution 106 over the set of possible action data elements. The set of possible action data elements may be any suitable set of data elements, for example, integer values in the range [0, 255], or a predefined set of set paddings. The system 100 then selects an action data element 104 using the score distribution 106 over the set of possible action data elements. For example, the system 100 may select the action data element 104 with the highest score according to the score distribution 106. As another example, the system 100 may sample action data elements from the set of possible action data elements according to a probability distribution over the set of possible action data elements, which may be generated by processing the score distribution 106 using, for example, a softmax function.

[0097] In some cases, for one or more positions in the sequence of action data elements defining the current action 102, the set of valid action data elements at that position may be a proper subset (i.e., less than all) of the set of possible action data elements. Action data elements at a position may be said to be “valid” if an action including the action data elements at that position represents a feasible action that may be performed by the agent. For example, if the action data elements at a position represent torques to be applied to joints of a robotic agent, the robotic agent may apply M possible torques to the joints, and the set of possible action data elements includes N>M action data elements, then M of the possible action data elements may be designated as being valid action data elements at that position. The system may ensure that the action data element selected at each position is a valid action data element, for example, by selecting the valid action data element with the highest score according to a score distribution over the set of possible action data elements at that position.

[0098] After generating each action data element 104, the system 100 updates the current task state 110 by concatenating the action data element 104 to the current task state 110 before generating the next action data element 104 in the sequence of action data elements 104 that defines the current action 102. Thus, the action selection neural network 108 autoregressively generates the sequence of action data elements 104, i.e., the action data element 104 at each position is generated by processing the current task state 110, which includes the action data elements 104 generated for each preceding position. An example of using the action selection neural network 108 to select an action to be performed by an agent by autoregressively generating action data elements is shown with reference to Figures 3A and 3B.

[0099] A sequence of action data elements 104 defines the actions 102 to be performed by the agent 120 at a time step. For example, if the agent is a mechanical agent, the action data element 104 at each position in the sequence may define the torque to be applied to a corresponding joint of the robot. As another example, if the agent is an autonomous vehicle, the action data element 104 at one position may define the acceleration / deceleration to be implemented by the vehicle, and the action data element 104 at another position may define the steering to be implemented by the vehicle.

[0100] In some cases, hyperparameters of system 100 may specify the maximum length of current task state 110, i.e., the maximum number of data elements that may be included in the current task state 110. As system 100 concatenates data elements representing new observations and actions onto the "terminal" end of the current task state, the length of the current task state increases. Thus, the system may remove data elements from the "initial" end of the current task state as necessary to ensure that the length of the current task state remains at most the maximum length. (The terminal end of the current task state refers to the position occupied by the last data element in the sequence of data elements representing the current task state, and the initial end of the current task state refers to the position occupied by the first data element in the sequence of data elements representing the current task state.)

[0101] The action selection neural network 108 may have any suitable neural network architecture that enables it to perform its described function, i.e., processing the current task state 110 and generating a score distribution over a set of possible action data elements. In particular, the action selection neural network may include any suitable neural network layers (e.g., attention layers, convolutional layers, fully connected layers, etc.) in any suitable number (e.g., 5, 10, or 100 layers) and connected in any suitable configuration (e.g., as a linear sequence of layers).

[0102] Several examples of possible architectures for the action selection neural network 108 are described next. In each of these examples, the action selection neural network may include an embedding layer configured to map each data element in the sequence of data elements defining the current task state 110, represented as a numerical value, to a corresponding embedding in an embedding space. The embedding layer may maintain, without modification, the data elements in the sequence of data elements defining the current task state 110 that are already represented as an embedding. That is, the embedding layer may represent the current task state 110 as a collection of embeddings, for example, by replacing each numerical value included in the current task state 110 with its corresponding embedding according to a predefined mapping from numerical values to embeddings.

[0103] In some cases, for each position in the current task state 110, the embedding layer can combine (e.g., add or average) the embedding for the position with a position embedding that represents the position in the current task state. Such position embedding can enable the action selection neural network to fully exploit the order of data elements in the current task state 110 without relying on recurrence or convolution.

[0104] In one example, the action selection neural network 108 can process the current task state 110 using an embedding layer to generate a collection of embeddings that represent the current task state 110. The action selection neural network 108 can then process the embeddings that represent the current task state 110 using a sequence of neural network layers, including one or more self-attention layers (e.g., self-attention layers that use a query-key-value attention mechanism), to generate a set of updated embeddings. The action selection neural network 108 can process the updated embeddings using one or more final neural network layers to project the updated embeddings onto score distributions over the set of possible action data elements. In a particular example, the action selection neural network 108 may have a Transformer Neural Network (a neural network characterized by having a series of self-attention neural network layers) architecture, e.g., a Transformer Neural Network decoder, as described with reference to A. Vaswani et al., "Attention is all you need," arXiv:1706.03762v5, December 6, 2017. The Transformer Neural Network may include memory to facilitate processing longer sequences of data elements representing the current state of the task. For example, the Transformer Neural Network may have a Transformer-XL architecture, as described in "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context," Dai et al., arXiv:1901.02860v3, June 2, 2019.

[0105] In another example, the action selection neural network 108 may include an embedding layer followed by a fully connected layer that is applied separately to each embedding representing each data element in the sequence of data elements representing the current task state. The updated embeddings produced by the fully connected layer may be combined (e.g., averaged) and then processed by a final fully connected neural network layer to generate a score distribution over the set of possible action data elements.

[0106] In another example, the action selection neural network 108 may be a recurrent neural network (RNN), e.g., a long short-term memory (LSTM) neural network. The RNN may be configured to process embeddings representing data elements to update its hidden state (e.g., cell states) and, in some cases, process the updated hidden state to generate a score distribution over a set of possible action data elements. After receiving an observation, the RNN may iteratively update its hidden state by processing each embedding corresponding to each data element in the sequence of data elements representing the observation, one by one, in order starting from the first position in the sequence. The RNN may then autoregressively generate a sequence of data elements defining an action to be performed in response to the observation. In particular, for each position in the sequence of action data elements, the RNN processes its current hidden state to generate a score distribution over the set of possible action data elements that is used to select an action data element for that position. The RNN then processes the embedding representing the action data element selected for that position to update its hidden state before generating a score distribution over the set of possible action data elements for the next position.

[0107] In some implementations, rather than receiving observations from the external environment, the action selection system 100 can simulate observations received from the environment. More specifically, at each time step, the action selection system 100 can autoregressively generate a sequence of data elements representing the current action performed by the agent at that time step, and then autoregressively generate a sequence of data elements representing the observations at the next time step.

[0108] The action selection system 100 can autoregressively generate a sequence of data elements representing an observation at the next time step by starting from the first position in the sequence of data elements representing the observation and successively generating a respective data element for each position. Specifically, for each position in the sequence of data elements representing the observation, the system can process the current task state 110 using an action selection neural network to generate a score distribution over the set of possible data elements. The system can select a data element for a position in the sequence of data elements representing the observation according to the score distribution. For example, the system can select the data element with the highest score under the score distribution, or the system can sample data elements according to a probability distribution over the set of data elements defined by the score distribution over the set of data elements. The system can then update the current task state 110 by concatenating the selected data element for a position in the sequence of data elements representing the observation to the current task state 110. The system can continue to autoregressively generate a sequence of data elements representing the observation until a termination criterion is met, for example, until the system has generated a predefined number of data elements that collectively define the observation.

[0109] In some implementations, action selection system 100 may be configured to generate actions only, i.e., without receiving or generating observations, and particularly by not including any observations (with the possible exception of any observations included in prompts provided to action selection system 100) in the current task state 110. For example, action selection system 100 may be configured to perform a text question-answering task (as described in more detail below) by generating a sequence of actions that represent text responses to a question (i.e., without generating or receiving any intervening observations).

[0110] In some implementations, action selection system 100 may be configured to generate observations only, i.e., without receiving or generating actions, and particularly by not including any actions in the current task state 110 (with the possible exception of any actions included in the prompt provided to action selection system 100). For example, action selection system 100 may be configured to perform a video generation task, particularly by generating a sequence of video frames, where each video frame represents a respective observation. For example, action selection system 100 may be configured to receive a prompt that defines a topic for the video (e.g., "Generate a video about how to change a tire on a car"), and in response, action selection system 100 can generate a corresponding sequence of video frames related to the topic.

[0111] In some implementations, action selection system 100 may be configured, at each time step, to receive actions for the time step, generate a representation of the actions as a sequence of data elements, and concatenate the sequence of data elements representing the actions to the current task state 110. Action selection system 100 may then use action selection neural network 108 (as described above) to autoregressively generate a sequence of data elements representing the next observation and then proceed to the next time step. That is, action selection system 100 may receive actions from an external source; i.e., rather than using action selection neural network 108 to generate actions, it may use action selection neural network 108 to simulate observations that would result from performing the actions. In some cases, one or more of the actions may be specified by a user, for example, via an application programming interface (API) made available by action selection system 100. In some cases, one or more of the actions may be selected in any suitable manner, for example, using an external action selection policy that is parameterized by an external neural network.

[0112] Thus, in some implementations, action selection system 100 can generate actions while receiving observations from an external source (e.g., the environment), while in other implementations, action selection system 100 can generate observations while receiving actions from an external source (e.g., a user). Generating observations while receiving actions from an external source may enable action selection system 100 to generate a sequence of observations that simulate the effect of performing several actions in the environment.

[0113] The training engine 112 can train the action selection neural network 108 on training data 114, which includes a set of training examples. Each training example is represented as a sequence of data elements, e.g., a sequence of numbers, a sequence of embeddings, or a sequence containing numbers in some positions and embeddings in other positions. Thus, the action selection neural network 108 can be trained offline in a supervised manner. Alternatively, or in combination, the action selection neural network 108 can be partially or fully trained using offline or online reinforcement learning.

[0114] To train the action selection neural network 108 on training examples represented as a sequence of data elements, the training engine 112 can generate a respective prediction for each of one or more data elements included in the training examples. To generate a prediction for a specified data element in the training examples, the training engine 112 can process a subsequence of data elements preceding the specified data element in the training examples (i.e., collectively representing the “current task state”) to generate a score distribution over the set of possible data elements. The training engine 112 can determine the gradient of a loss function that measures the error, e.g., cross-entropy error, between (i) the score distribution over the set of possible data elements and (ii) the specified data element in the training examples. The training engine 112 can determine the gradient of the loss function with respect to the parameter values of the action selection neural network, for example, using backpropagation. The training engine 112 can use the gradient of the loss function to adjust the current values of the action selection neural network parameters using any suitable gradient descent optimization algorithm, for example, Adam or RMSprop.

[0115] Each data element included in the training examples may be designated as being an action data element, an observation data element, or a prompt data element. An action data element refers to a data element in a sequence of one or more data elements that represents an action (as described above). An observation data element refers to a data element in a sequence of one or more data elements that represents an observation. (In some cases, if the training engine 112 performs revenue adjustment, one of the observation data elements in a sequence of observation data elements for an observation can represent revenue, as described in more detail below.) A prompt data element refers to a data element in a sequence of one or more data elements that represents a prompt.

[0116] In some implementations, the training engine 112 trains the action selection neural network to predict only the actions included in each training example, i.e., the training engine 112 trains the action selection neural network to generate only the data elements in each training example that are designated as being action data elements (e.g., by masking other data elements).

[0117] In other implementations, the training engine 112 trains the action selection neural network to predict both the actions and the observations contained in each training example, i.e., the training engine 112 trains the action selection neural network to generate data elements in each training example that are designated as being either action data elements or observation data elements.

[0118] In general, the action selection neural network can generate a score distribution over the set of possible observation data elements in the same manner as the action selection neural network generates a score distribution over the set of possible action data elements. In some cases (e.g., when the set of possible observation data elements is different from the set of possible action data elements), the action selection neural network includes one output head (i.e., sub-network) configured to generate a score distribution over the set of possible observation data elements and a separate output head configured to generate a score distribution over the set of possible action data elements.

[0119] By training an action selection neural network to predict both the observations and actions contained in the training examples, the action selection neural network can implicitly learn a model of the environmental dynamics, thereby enabling the action selection neural network to more effectively select actions for performing a task.

[0120] When an action selection neural network is used to select an action to be performed to control an agent, the action selection neural network autoregressively generates action data elements. However, it will be appreciated that during training, the training engine can use the action selection neural network to generate predictions for each data element included in each training example in parallel, thereby significantly increasing the efficiency of training.

[0121] In some cases, the training engine can autoregressively generate predictions for actions, observations, or both during training. For example, to generate predictions for action data elements that define actions in the training examples, the training engine can initialize a current task state that includes a subsequence of data elements prior to a first action data element in the sequence of action data elements that define actions in the training examples. The training engine can then process the current task state as described above to autoregressively generate predictions for action data elements that define actions.

[0122] In general, any sequence of data elements from any suitable source may be used as training examples for training the action selection neural network 108. Thus, the training engine 112 may train the action selection neural network 108 on training examples from a wide variety of sources, including, for example, both simulated data and real-world data. Several exemplary techniques for generating training examples for training the action selection neural network are described in more detail below.

[0123] In one example, the training engine 112 can generate training examples that represent the interaction of an agent with an environment over a sequence of time steps. The interaction of an agent with an environment may be represented in the following form:

[0124]

number

[0125] where N is the number of time steps and s i is the state of the environment at time step i, and a i is the action taken by the agent at time step i, and ri is the reward received at time step i. In general, for each state s i and each action a i may be represented as an ordered collection of numbers, e.g., a vector, matrix, or other tensor of numbers. (The number of numbers in an ordered collection of numbers defining an action is sometimes called the "dimensionality" of the action.) The use of rewards is optional. Rewards may be used, for example, to filter training examples to select those with at least a threshold percentage of the revenue achieved by the expert agent performing the task.

[0126] To generate training examples representing the agent's interactions with the environment, the training engine 112 represents each observation as a sequence of data elements, e.g., a sequence of numerical values or a sequence of embeddings. For example, the training engine 112 can represent each collection of numerical values defining each observation as a sequence of numerical values, e.g., by concatenating the numerical values defining the observation into a sequence of numerical values in an arbitrary but fixed order. (An exemplary technique for representing observations as a sequence of embeddings is described above.) The training engine 112 can perform profit adjustment by determining a profit, e.g., by calculating a time-discounted reward sum, and then concatenating the profit to the sequence of data elements representing each observation. In some cases, the training engine 112 can normalize the profit, e.g., by dividing the profit by the expert profit for the task being performed by the agent, e.g., the average profit value achieved when the agent performed the task one or more times under the control of an expert, e.g., a human expert. The training engine 112 also represents each action as a sequence of data elements; for example, the training engine 112 may represent an action as a sequence of numbers by concatenating collections of numbers representing the action into a sequence of numbers. The training engine 112 then concatenates each sequence of data elements representing each observation and each action at each time step into one sequence of data elements. As part of generating the training examples, the training engine 112 may, in some cases, discretize the numbers in the collections of numbers representing the observations, the actions, or both.

[0127] In some cases, the training engine 112 may generate a representation of a prompt for a training example as a sequence of data elements and concatenate the prompt to the sequence of data elements that represent the training example.

[0128] Performing revenue adjustment enables the training engine 112 to effectively train the action selection neural network 108 on training examples in which the agent receives a variety of possible revenues, including low revenues. A low revenue associated with a training example may indicate that the training example represents an agent interaction with the environment during which the agent failed to perform its task. In the absence of revenue adjustment, training the action selection neural network 108 on training examples associated with low revenues may degrade the performance of the action selection neural network 108, for example, by reinforcing an ineffective action selection policy represented by the training example. Performing revenue adjustment enables the action selection neural network to distinguish between training examples representing effective and ineffective action selection policies, thereby improving the performance of the action selection neural network 108.

[0129] In another example, the training engine 112 can generate "language modeling" training examples that represent sequences of text in a natural language. The training engine 112 can represent the sequence of text as a sequence of tokens from a predefined set of possible tokens (e.g., characters, n-grams, or words), and then replace each token with a corresponding data element, e.g., an integer identifier that indexes the token in the set of possible tokens, or an embedding. The training examples can then be represented by a sequence of data elements that identify the sequence of tokens, where each token is designated as being an action data element (i.e., in this case, the training examples would not include any observation data elements).

[0130] In another example, the training engine 112 can generate "image captioning" training examples that represent (i) an image and (ii) an image caption that defines a sequence of text that describes the content of the image. For example, the training engine 112 can generate the training examples by concatenating respective sequences of data elements that represent the image and the image caption. The sequence of data elements that represent the image can be designated as a sequence of observation data elements, and the sequence of data elements that represent the image caption can be designated as a sequence of action data elements.

[0131] In another example, the training engine 112 can generate "text question-answer" training examples that represent (i) a text question and (ii) a text response that responds to the text question. For example, the training engine 112 can generate the training examples by concatenating respective sequences of data elements that represent the text question and the text response. The sequence of data elements that represent the text question can be specified as a sequence of observation data elements, and the sequence of data elements that represent the text response can be specified as a sequence of action data elements.

[0132] In another example, the training engine 112 can generate "visual question-answer" training examples that represent (i) an image and a text question associated with the image, and (ii) a text response that responds to the text question. For example, the training engine 112 can generate the training examples by concatenating respective sequences of data elements that represent the image, the text question, and the text response. The sequence of data elements that represent the image and the text question can be designated as a sequence of observation data elements, and the sequence of data elements that represent the text response can be designated as a sequence of action data elements.

[0133] In another example, the training engine 112 can generate "image classification" training examples that represent (i) an image and (ii) a classification of the image into a category from a predefined set of categories. For example, each category can represent a respective type of object, and an image can be classified as being in that category if it shows an object of the type represented by the category, and each category can be represented by a respective numeric value. The training engine 112 can generate the training examples by concatenating (i) a sequence of data elements that represent the image and (ii) a numeric value that represents the classification of the image. The sequence of data elements that represent the image can be designated as a sequence of observation data elements, and the numeric value that represents the classification of the image can be designated as an action data element.

[0134] The training engine 112 can train the action selection neural network 108 on training examples from a diverse set of multiple different domains. In particular, the training engine 112 can train the action selection neural network on training examples representing interactions between multiple different agents and multiple different environments to perform multiple different tasks. (Examples of possible agents, environments, and tasks are described above.) Training the action selection neural network 108 on training examples from multiple domains can encode a flexible and transferable understanding of agent control in the parameters of the action selection neural network, thereby enabling the action selection neural network to generalize quickly and effectively to new domains. In particular, training the action selection neural network 108 on multiple domains can enable the action selection neural network 108 to achieve acceptable performance on tasks in new domains after being trained on a small number of training examples from the new domain. In some cases, training in multiple domains can enable the action selection neural network 108 to achieve acceptable performance on a task in a new domain even if the action selection neural network has not been trained on any training examples from the new domain.

[0135] Training the action selection neural network 108 on additional training examples other than those representing agent interaction with the environment, such as the language modeling, image captioning, text question answering, visual question answering, and image classification training examples described above, can accelerate training and improve the performance of the action selection neural network. For example, training the action selection neural network on language modeling training examples can improve the ability of the action selection neural network to implicitly infer the meaning of natural language prompts provided for control tasks. This can, for example, facilitate generalization to tasks in environments for which the system was not specifically trained.

[0136] In general, training examples from different domains can use sequences of data elements of different lengths to represent actions and observations. For example, as shown in FIG. 2, training examples from “Domain #1” 202 use sequences of four data elements to represent observations and sequences of two data elements to represent actions, while training examples from “Domain #2” 204 use sequences of three data elements to represent observations and sequences of three data elements to represent actions. This may present a problem for conventional action selection neural networks, which have neural network architectures configured to process fixed-size observations and generate fixed-size actions, for example. In contrast, the operation of the action selection neural network 108 can be flexibly adapted to process training examples from any domain, regardless of the domain-specific dimensionality of the observations and actions. For example, to generate actions with a dimensionality appropriate for a particular domain, the action selection neural network 108 can continue to autoregressively sample action data elements until the generated actions have the appropriate dimensionality.

[0137] 3A and 3B illustrate operations performed by the action selection system 100 to select actions to be performed by an agent that interacts with an environment to accomplish a task.

[0138] 3A illustrates operations performed to autoregressively generate a sequence of action data elements representing actions to be performed by an agent at the first time step, i.e., T=0. The system, in this example, initializes a current task state 304 using a prompt 302. The prompt 302, represented as a sequence of data elements, may include any suitable data related to the task to be performed by the agent, such as a demonstration of the task or natural language instructions related to the task.

[0139] The system 100 receives observations from the environment, representing the current state of the environment, for example, in the form of an image of the environment. The system represents the observations as a sequence of observation data elements 310, for example, by concatenating numbers in a collection of numbers representing the observations into a sequence in an arbitrary but fixed order.

[0140] The system 100 then links the observation data element to the current task state 304 .

[0141] The system 100 uses the action selection neural network 108 to process the current task state 304 to generate a probability distribution over the set of possible action data elements, and then selects an action data element 312 according to the probability distribution over the set of possible action data elements.

[0142] The system 100 concatenates the action data element 312 to the current task state 306 and uses the action selection neural network 108 to process the updated task state 306 to generate another action data element 314. More specifically, the system processes the updated task state 306 to generate a probability distribution over the set of possible action data elements, and then selects the action data element 314 according to the probability distribution over the set of possible action data elements.

[0143] The generated sequence of action data elements (ie, including action data elements 312 and 314) defines an action 316 to be performed by the agent in the first time step.

[0144] FIG. 3B illustrates the operations performed to autoregressively generate a sequence of action data elements representing actions to be performed at the second time step, ie, T=1.

[0145] The agent performs a selected action 316 at a first time step (T=0), and the environment transitions to a new state as a result of the action performed by the agent. The system receives observations characterizing the new state of the environment at a second time step, represents the observations as a sequence of observation data elements 318, and concatenates the observation data elements to a current task state 326. Thus, the current task state 326 includes the prompt 302, a sequence of observation data elements representing the observations at the first time step, a sequence of action data elements representing actions performed by the agent at the first time step, and a sequence of observation data elements 318 representing the observations at the second time step.

[0146] The system 100 uses an action selection neural network to process the current task state 326 to generate a score distribution over the set of possible action data elements and select an action data element 320 according to the score distribution over the set of possible action data elements.

[0147] The system then links the action data element 320 to the current task state 328 .

[0148] The system processes the updated task state 328 using the action selection neural network 108 to generate another score distribution over the set of possible action data elements, and selects an action data element 322 according to the score distribution over the set of possible action data elements. The system concatenates the action data element 322 with the current task state 330 to provide the current task state 330 for use in selecting an action to be performed by the agent in the next time step.

[0149] The generated sequence of action data elements 320 and 322 defines the action 324 to be performed by the agent in the second time step.

[0150] 4 is a flow diagram of an exemplary process 400 for selecting an action to be performed by an agent to interact with an environment at a current time step. For convenience, process 400 is described as being performed by one or more computer systems located at one or more locations. For example, an action selection system suitably programmed in accordance with this specification, such as action selection system 100 of FIG. 1, can perform process 400.

[0151] The system generates 402 a current representation of the state of the task being performed by the agent in the environment as of the current time step as a sequence of data elements, which may be, for example, a sequence of numbers, a sequence of embeddings, or a sequence containing numbers in some positions and embeddings in other positions.

[0152] The system autoregressively generates a sequence of data elements representing a current action to be performed by the agent at the current time step. Specifically, the system performs steps 404-410 for each position in the sequence of data elements representing the current action, starting from the first position. For convenience, steps 404-410 are described as being performed for the current position in the sequence of data elements representing the current action.

[0153] The system uses an action selection neural network to process the current representation of the state of the task to generate a score distribution over the set of possible data elements (404).

[0154] The system selects (406) a data element for the current position in the sequence of data elements representing the current action according to the score distribution.

[0155] The system updates 408 the current representation of the task's state by concatenating the selected data element for the position to the current representation of the task's state.

[0156] The system determines whether the current action is complete (410). If the current position is the last position in the sequence of data elements representing the current action, the system determines that the current action is complete and proceeds to step 412. Otherwise, the system determines that the current action is not complete and loops back to step 404.

[0157] After autoregressively generating a sequence of data elements representing the current action, the system causes the agent to perform the current action at the current time step (412).

[0158] This specification uses the term "configured" in connection with systems 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 software, firmware, hardware, or a combination thereof that causes the system to perform the operation or action during operation. 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 the operation or action.

[0159] Embodiments of 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, or a combination of one or more of them, including the structures disclosed herein and their structural equivalents. Embodiments of the subject matter described herein may be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium for execution by or to control the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or sequential access memory device, or a combination of one or more of them. Alternatively or additionally, 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 a suitable receiver apparatus for execution by a data processing apparatus.

[0160] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or computers. An apparatus may also be or further include special-purpose logic circuitry, e.g., an FPGA (field-programmable gate array) or an ASIC (application-specific integrated circuit). In some cases, an apparatus may include, in addition to hardware, code that creates an execution environment for a computer program, 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 these.

[0161] A computer program, which may also be called or described as a program, software, software application, app, module, software module, script, or code, may be written in any type of programming language, including compiled or interpreted, or declarative or procedural, and may 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 may 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 storing one or more modules, subprograms, or portions of code. A computer program may be deployed to be executed on one computer or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communications network.

[0162] The term "engine" is used broadly herein 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 installed on one or more computers at one or more locations. In some cases, one or more computers will be dedicated to a particular engine, and in other cases, multiple engines may be installed and running on the same one or more computers.

[0163] 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 operating on input data and generating output. The processes and logic flows may also be performed by special purpose logic circuitry, e.g., an FPGA or ASIC, or a combination of special purpose logic circuitry and one or more programmed computers.

[0164] A computer suitable for running a computer program may be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. Generally, the 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 memory may be supplemented by, or incorporated in, special-purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to, one or more mass storage devices, e.g., magnetic, magneto-optical, or optical disks, for storing data therefrom, or for transferring data thereto, or both. However, a computer need not have such devices. Moreover, a computer may be embedded in another device, e.g., a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0165] 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.

[0166] To provide for user interaction, embodiments of the subject matter described herein may 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 pointing device, e.g., a mouse or trackball, by which the user can provide input to the computer. Other types of devices may similarly be used to provide for user interaction; for example, feedback provided to the user may be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback, and input from the user may be received in any form, including acoustic, voice, or tactile input. Additionally, a computer may interact with a user by sending documents to and receiving documents from a device used by the user, e.g., by sending a web page to a web browser on the user's device in response to a request received from the web browser. A computer may also interact with a user by sending text messages or other forms of messages to a personal device, e.g., a smartphone running a messaging application, and receiving a response message from the user in return.

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

[0168] The machine learning model may be implemented and deployed using a machine learning framework, for example, the TensorFlow framework.

[0169] Embodiments of the subject matter described herein may be implemented in a computing system that includes back-end components, e.g., as data servers, or middleware components, e.g., application servers, or front-end components, e.g., client computers having graphical user interfaces, web browsers, or apps through which users can interact with implementations of the subject matter described herein, or any combination of one or more such back-end, middleware, or front-end components. The components of the system may be interconnected by any form or medium of digital data communication, e.g., a communications network. Examples of communications networks include local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

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

[0171] While the specification contains numerous specific implementation details, these should not be construed as limiting the scope of any invention or of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Some features described herein 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. Moreover, while features may be described above as operative in some combinations, and even initially claimed as such, one or more features from a claimed combination may, in some cases, be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of a subcombination.

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

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

[0174] 100 Action Selection System, System 102 Action, Current Action 104, 312, 320, 322 Action Data Elements 106 Score Distribution 108 Action Selection Neural Network 110, 304, 326, 330 Current task state 112 Training Engine 114 training data 116 observations, current observations 118 Environment 120 Agents, Robot Agents 302 prompt 306, 328 Current task state, Updated task state 310, 318 Observation Data Elements 314 Another Action Data Element, Action Data Element 316, 324 Action

Claims

1. 1. A method implemented by one or more computers for selecting actions to be performed by an agent to interact with an environment using an action selection neural network, comprising: At each time step in the sequence of time steps, generating a current representation of the state of the task being performed by said agent in said environment as of a current time step as a sequence of data elements; autoregressively generating a sequence of data elements representing a current action to be performed by the agent at the current time step; Starting from the first position in the sequence of data elements representing the current action, for each position: processing the current representation of the state of the task using the action selection neural network to generate a score distribution over a set of possible data elements; selecting data elements according to the score distribution, wherein the selecting comprises at least one of selecting a data element having a highest score under the score distribution or sampling data elements according to a probability distribution over the set of possible data elements defined by the score distribution over the set of possible data elements; updating the current representation of the state of the task by concatenating the selected data element to the current representation of the state of the task; and autoregressively generating the sequence of data elements representing the current action, and then causing the agent to perform the current action at the current time step; A method comprising:

2. generating, for each time step in the sequence of time steps, the current representation of the state of the task as of the current time step, receiving a current observation characterizing a state of the environment at the current time step; generating a representation of the current observation as a sequence of data elements; including the representation of the current observation as a sequence of data elements in the current representation of the state of the task as of the current time step; 2. The method of claim 1, comprising:

3. the current observation is defined by a set of numerical values; generating the representation of the current observation as a sequence of data elements, 3. The method of claim 2, comprising concatenating each number in the collection of numbers that defines the current observation into a sequence of numbers in a predefined order.

4. generating the representation of the current observation as a sequence of data elements, The method of claim 3 , further comprising discretizing each number in the collection of numbers that defines the current observation.

5. The method of claim 2 , wherein the current observation characterizing the current state of the environment at the current time step comprises an image defined by an array of pixels.

6. generating the representation of the current observation as a sequence of data elements, 3. The method of claim 2, comprising combining a target return to be achieved by an interaction of the agent with the environment with the representation of the current observation as a sequence of data elements, the target return defining a cumulative measure of reward to be achieved as a result of the interaction of the agent with the environment.

7. for each time step after a first time step in the sequence of time steps, including the representation of the current observation as a sequence of data elements in the current representation of the state of the task as of the current time step; receiving a representation of the state of the task at a previous time step as a sequence of data elements; concatenating the representation of the current observation as a sequence of data elements with the representation of the state of the task as of the previous time step as a sequence of data elements to generate the current representation of the state of the task as of the current time step; 3. The method of claim 2, comprising:

8. 8. The method of claim 7, wherein the representation of the state of the task at the previous time step represents, for each time step preceding the current time step, (i) a respective observation characterizing the state of the environment at that time step, and (ii) a respective action performed by the agent at that time step.

9. at a first time step in the sequence of time steps, including the representation of the current observation as a sequence of data elements in the current representation of the state of the task as of the current time step; receiving a prompt comprising data characterizing the task to be performed by the agent in the environment; generating a representation of the prompt as a sequence of data elements; concatenating the representation of the current observation as a sequence of data elements with the representation of the prompt as a sequence of data elements to generate the current representation of the state of the task as of the current time step; 3. The method of claim 2, comprising:

10. 10. The method of claim 9, wherein the prompt comprises one or more of a demonstration of the task, a goal observation that characterizes a goal state of the environment, or a sequence of text in natural language that provides instructions related to the task.

11. the action selection neural network is trained on a set of training examples; For each training example, the training examples are represented as sequences of data elements; at least one of the data elements in the sequence of data elements representing the training examples is designated as an action data element; 2. The method of claim 1, wherein training the action selection neural network on the training examples comprises training the action selection neural network to generate the action data elements included in the training examples.

12. the set of training examples includes respective training examples from a plurality of different control domains; Each control domain is associated with (i) a corresponding agent, (ii) a corresponding environment, and (iii) a corresponding task; The method of claim 11 , wherein each training example from each control domain characterizes an interaction between the corresponding agent and the corresponding environment by performing an action to accomplish the corresponding task.

13. 13. The method of claim 12, wherein the plurality of different control domains comprises a first control domain in which the corresponding environmental observations have a first dimensionality and a second control domain in which the corresponding environmental observations have a second, different dimensionality.

14. 13. The method of claim 12, wherein the plurality of different control domains comprises a first control domain in which actions performed by the corresponding agents have a first dimensionality, and a second control domain in which actions performed by the corresponding agents have a second, different dimensionality.

15. the set of training examples includes a plurality of language modeling training examples; The method of claim 11 , wherein each language modeling training example represents a sequence of text in a natural language.

16. The method of claim 1 , wherein the action selection neural network comprises multiple self-attention neural network layers.

17. 2. The method of claim 1, wherein, for each time step in the sequence of time steps, the sequence of data elements representing the state of the task as of the current time step comprises a sequence of numeric values, a sequence of paddings, or a sequence containing numeric values in some positions and paddings in other positions.

18. the current observation comprises an image; generating the representation of the current observation as a sequence of data elements, generating a respective initial patch embedding corresponding to each of a plurality of patches in the image; processing the initial patch embeddings using an encoder neural network to generate a respective final patch embedding for each of the plurality of patches in the image; Including, The method of claim 2 , wherein each final patch embedding is included as a respective data element in the sequence of data elements representing the current observation.

19. generating respective initial patch embeddings corresponding to patches in the image, generating pixel embeddings representing pixels in the patch in the image; generating a patch position embedding representing the position of the patch in the image; generating the initial patch embedding for the patch by combining the pixel embedding and the patch position embedding for the patch; 20. The method of claim 18, comprising:

20. 20. The method of claim 18, wherein the encoder neural network comprises one or more self-attention neural network layers.

21. The method of claim 18 , wherein the encoder neural network comprises one or more residual blocks.

22. The method of claim 1 , wherein the agent is a mechanical agent that interacts with a real-world environment.

23. 23. The method of claim 22, wherein selecting an action to be performed by the mechanical agent comprises selecting an action that causes the mechanical agent to physically manipulate one or more objects in the environment.

24. 1. A system comprising: one or more computers; one or more storage devices communicatively coupled to the one or more computers; wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of the respective method of any one of claims 1 to 23.

25. 24. One or more non-transitory computer-readable storage media having instructions stored thereon that, when executed by one or more computers, cause the one or more computers to perform the operations of the respective method of any one of claims 1 to 23.

Images