Controlling Agents through In-Context Reinforcement Learning

Abstract

A method, system, and apparatus for controlling an agent, including a computer program encoded on a computer storage medium, wherein, in particular, the agent can be controlled using an action selection neural network that performs in-context reinforcement learning when controlling the agent in new tasks.

Description

[Technical Field]

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

[0002] This specification relates to processing data using machine learning models.

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

[0004] Some machine learning models are deep models that use 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, where each layer applies a nonlinear transformation to the received input to generate an output. Summary of the Invention

[0005] This specification generally describes a system, implemented as a computer program on one or more computers in one or more locations, that controls an agent that interacts with an environment and performs tasks within the environment using an action selection neural network.

[0006] More specifically, the agent selection neural network represents an "in-context" reinforcement learning algorithm. That is, because the amount of contextual data has increased while the agent was controlled using the action selection neural network, by conditioning on contextual data from prior interactions with the environment, the action selection neural network may select actions that result in improved task performance (compared to previous points in time while controlling the agent) without updating the parameters of the action selection neural network. In other words, the action selection neural network can "mimic" the performance of a reinforcement learning algorithm as the amount of available data increases, without the need to update the parameters of the neural network, i.e., without the need to further train the neural network.

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

[0008] This specification describes an action selection neural network that can be trained on a training dataset and then used to perform in-context reinforcement learning on new tasks. When performing new tasks, the action selection neural network can be significantly more data-efficient than the reinforcement learning algorithm(s) used to generate the training dataset. This is because, for example, a multi-actor algorithm is "distilled" into a single-actor algorithm, the system subsamples training episodes when generating the training dataset, or both.

[0009] Furthermore, "learning" a new task using an action selection neural network also consumes significantly fewer computational resources than learning a policy for a new task using an RL algorithm because the action selection neural network no longer needs to be updated while "learning," eliminating the need for a computationally expensive backward pass. That is, the action selection neural network improves on the new task solely by observing more context, without updating the neural network weights, so no gradient calculations and therefore no backward pass are required, while still improving performance on the new task.

[0010] Furthermore, distributed reinforcement learning algorithms require a large amount of network communication between distributed actors and learners during the learning process. This network communication can be significantly reduced or even eliminated by utilizing an already trained action-selection neural network. That is, by utilizing the described "in-context" reinforcement learning scheme, a system can achieve performance for a given new task with a single actor implemented on a single set of one or more hardware devices that is comparable to or better than learning the task using a distributed reinforcement learning algorithm that requires multiple actors and one or more learners, each implemented on a different set of one or more hardware devices. Thus, network communication is significantly reduced because weight updates do not need to be sent between actors and learners and transitions are not required to be sampled from a replay buffer.

[0011] Furthermore, after training, when "learning in-context" new tasks, the storage requirements are at least comparable to, and often significantly reduced compared to, distributed reinforcement learning algorithms. This is because only the tokenized observations, actions, and rewards needed for the current context of the action-selection neural network at the current time step need to be stored, and any tokenized observations and rewards that are no longer needed for the context can be discarded. In contrast, distributed reinforcement learning algorithms are required to maintain a replay buffer containing a large number of transitions that the learner can sample to train the neural network. By eliminating the requirement to maintain this large replay buffer, the described technique significantly reduces the storage requirements for "learning" new tasks.

[0012] Furthermore, the action selection neural network can be trained purely "offline," i.e., without needing to be used to control an agent, which not only avoids damage or wear to the real-world agent but also allows the training workload to be amortized and parallelized.

[0013] Furthermore, once trained, the same action selection neural network can be used to "learn" many new tasks without the need for any further training, greatly reducing the amount of computational resources required within a multitasking system.

[0014] 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]

[0015] [Figure 1] 1 illustrates an exemplary action selection system. [Figure 2]1 is a flow diagram of an example process for training an action selection neural network. [Figure 3] 1 is a flow diagram of an example process for generating a training history sequence for a given task. [Figure 4] FIG. 1 is a diagram of an example of training an action selection neural network. [Figure 5] FIG. 1 is a flow diagram of an exemplary process for controlling an agent using an action selection neural network. [Figure 6] 1 shows an example of an attention map generated by an action selection neural network. [Figure 7] We show the performance of the described technique compared to a baseline technique. DETAILED DESCRIPTION OF THE INVENTION

[0016] Like reference symbols and designations in the various drawings refer to like elements.

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

[0018] The action selection system 100 controls an agent 104 that interacts with an environment 106 to accomplish a task by selecting an action 108 to be performed by the agent 104 at each of multiple time steps during the execution of an episode of the task.

[0019] As general examples, tasks may include, for example, one or more of: navigating to a specified location in an environment, identifying a particular object in an environment, manipulating a particular object in a specified manner, controlling an item of equipment to meet a criterion, distributing resources among devices, etc. More generally, tasks are specified in terms of the reward received, i.e., such that an episodic return is maximized upon successful completion of the task. Rewards and returns are described in more detail below. Example agents, tasks, and environments are also provided below.

[0020] A task "episode" is a sequence of interactions in which an agent attempts to perform a single instance of a task, starting from some starting state of the environment. In other words, each task episode begins with the environment in an initial state, e.g., a fixed initial state or a randomly selected initial state, and ends when the agent successfully completes the task or when some termination criterion is met, e.g., the environment enters a state designated as a terminal state or the agent performs a threshold number of actions without successfully completing the task.

[0021] At each time step of any given task episode, the system 100 receives observations 110 that characterize the current state of the environment 106 at the time step, and in response selects an action 108 to be performed by the agent 104 at the time step. After the agent performs the action 108, the environment 106 transitions to a new state, and the system 100 receives a reward 130 from the environment 106.

[0022] Generally, the reward 130 is a scalar numerical value that characterizes the progress of the agent 104 toward completing the task.

[0023] As a particular example, reward 130 may be a sparse binary reward that is zero if the action being performed does not result in successful completion of the task, i.e., is only non-zero, e.g., equal to 1, if the action being performed results in successful completion of the task.

[0024] As another particular example, reward 130 may be a dense reward that measures the agent's progress toward completing the task from individual observations received during an episode of attempting to perform the task, i.e., a non-zero reward can and will be received frequently before the task is successfully completed.

[0025] While performing any given task episode, the system 100 selects actions to attempt to maximize the returns received over the course of the task episode.

[0026] That is, at each time step in the episode, the system 100 selects an action that attempts to maximize the return received for the remainder of the task episode starting at that time step.

[0027] In general, at any given time step, the returns received are a combination of the returns received at time steps following the given time step in the episode.

[0028] For example, at time step t, the return is

number

[0029] To control the agent, at each time step of an episode, the action selection subsystem 102 of the system 100 uses an action selection neural network 102 to select an action 108 to be performed by the agent 104 at the time step.

[0030] At a high level, the system 100 generates an input sequence 112 from observations 110 and then uses an action selection neural network 102 to process the input sequence 112 and generate a policy output for a time step.

[0031] As will be explained in more detail below, the input sequence 112 generally represents the observations 110 and further includes context data from previous time steps within the current episode and in previous episodes performing the task.

[0032] The system then uses the policy output to select an action 108 and causes the agent 104 to perform the selected action 108 .

[0033] More specifically, the agent selection neural network 102 represents an "in-context" reinforcement learning algorithm.

[0034] That is, by being conditioned on contextual data from prior interactions with the environment while the agent 104 was being controlled using the action selection neural network 102, and because the amount of contextual data available has increased, the action selection neural network 102 may select actions that result in improved task performance (compared to previous points while controlling the agent) without updating the (learnable) parameters, e.g., weights, of the action selection neural network 102.

[0035] In other words, the action selection neural network 102 can "mimic" the performance of a reinforcement learning algorithm as the amount of available data increases, without the need to update the parameters of the neural network 102, i.e., without the need to further train the neural network.

[0036] In particular, the system 100 trains the action selection neural network 102 and then, after training, uses the action selection neural network 120 as an “in-context” reinforcement learning algorithm while performing a new task, i.e., performing the new task without the need to further train the action selection neural network 102.

[0037] For each new task, performance improves as the amount of available contextual data increases, for example, in terms of the average return obtained during any given task episode of the action selection neural network 102 in controlling the agent of the task, and without the need to train the action selection neural network 102.

[0038] To train the action selection neural network 102, the system obtains a training data set 150 that includes a respective training history sequence 152 for each of a number of tasks.

[0039] Generally, each training history sequence 152 includes a sequence of tokens representing transitions from multiple task episodes executed while training a policy, particularly an action selection policy, for a task through reinforcement learning. The policy for a task may be any suitable policy, e.g., any suitable policy that can be adjusted through training to improve the performance of the policy for the task, and may be represented by a neural network or other machine learning model having the same or different architecture as neural network 102. That is, any suitable training history sequence 152 for training any suitable policy may be incorporated into training dataset 150. Generally, an action selection policy may be a mechanism for selecting an action in response to an observation (of the environment).

[0040] As mentioned above, an "episode" of a task (or "task episode") is a sequence of interactions during which an agent attempts to perform an instance of a task, starting from some starting state of the environment.

[0041] In other words, each task episode begins with the environment in an initial state, e.g., a fixed initial state or a randomly selected initial state, and ends when the agent successfully completes the task or when some termination criterion is met, e.g., when the environment enters a state designated as a terminal state or the agent performs a threshold number of actions without successfully completing the task. At each time step of a task episode, the agent receives observations that characterize the state of the environment at that time step, performs actions in response to the observations, and receives a reward.

[0042] In other words, a task episode results in a sequence of "transitions," each corresponding to a different time step and each including a respective observation, a respective action, and a respective reward received in response to the action being performed.

[0043] The sequence of tokens in training history sequence 152 includes a respective episode subsequence for each of multiple task episodes that occurred during training of the task's policy.

[0044] Each episode subsequence for each task episode includes, for each transition from the task episode, a respective transition subsequence that includes (i) one or more tokens representing an observation at the transition, (ii) one or more tokens representing an action at the transition, and (iii) one or more tokens representing a reward at the transition.

[0045] As used herein, a "token" is a vector of numbers having a specified dimension.

[0046] In other words, tokenization system 120 in system 100 or a different tokenization system preprocesses the observations, actions, and rewards to "tokenize" them so that each observation, action, and reward is represented as one or more tokens, each with specified dimensions.

[0047] System 120 may represent the observations, actions, and rewards as tokens in any suitable manner.

[0048] One example of a technique for this is described in Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 2021.

[0049] Other example techniques for doing so are described in Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist Agent, 2022, arXiv:2205.06175.

[0050] That is, the system 120 may use one of the above techniques or other techniques to represent any given received action, observation, and reward as a respective set of one or more tokens having predetermined dimensionality.

[0051] The system 120 can generally use the same tokenization techniques when generating the data in the training sequences and when generating the input sequences 112 that are processed by the action selection neural network 102 after training.

[0052] The system 100 then trains the action selection neural network 102 based on the training data set. The result of the training is that the action selection neural network 102 is capable of performing "in-context" reinforcement learning as described above and below.

[0053] This training is described in more detail below with reference to FIGS.

[0054] After training, the system 100 can control the agent 104 to perform a task by controlling the agent 104 using the action selection neural network 102 at each of a number of time steps in a sequence of time steps in the "current" task episode.

[0055] At each time step, the system 100 receives a current observation 110 that characterizes the state of the environment 106 at the time step.

[0056] The tokenization system 120 then generates an input sequence 112 of tokens.

[0057] The input sequence 112 generally includes (i) one or more tokens representing the current observation 110, (ii) a respective current transition subsequence for each of one or more current episode transitions, and (iii) a respective previous transition subsequence for each of one or more previous episode transitions.

[0058] In some cases, tokens of the input sequence 112 other than the one or more tokens representing the current observation 110 have already been generated, and the tokenization system 120 need only tokenize the current observation 110 to generate one or more tokens representing the current observation 110.

[0059] Each current episode transition corresponds to a respective previous time step in the current task episode, and each transition subsequence of the current episode transition includes (a) one or more tokens representing observations received in the previous time step, (b) one or more tokens representing actions taken by the agent in response to the observations received in the previous time step, and (c) one or more tokens representing rewards received by the agent in response to taking the actions.

[0060] Each preceding episode transition corresponds to a respective previous time step in a respective preceding task episode (of the current task) performed by the agent prior to the current task episode, and each preceding transition subsequence of a preceding episode transition includes (a) one or more tokens representing observations received in the previous time step, (b) one or more tokens representing actions performed by the agent in response to the observations received in the previous time step, and (c) one or more tokens representing rewards received by the agent in response to performing the actions.

[0061] That is, the input sequence 112 provides context from earlier in the current episode as well as from previous task episodes.

[0062] As described above, the system 100 then processes the input sequence of tokens 112 using an action selection neural network to generate a policy output for the time step, uses the policy output to select an action, and causes the agent to perform the selected action.

[0063] The policy output may be any suitable output that defines a probability distribution over the set of actions. For example, the policy output may include a respective probability for each action in the set of actions. As another example, the policy output may include parameters of the probability distribution over the set of actions.

[0064] The system 100 can select an action 108, for example, by selecting the action with the highest probability according to the probability output or by sampling an action from a probability distribution defined by the policy output.

[0065] The action selection neural network 102 may be any suitable sequence model and may have, for example, any suitable architecture that enables the neural network 102 to map an input sequence of tokens to a probability distribution.

[0066] As an example, the action selection neural network 102 may be a causal transformer neural network, i.e., a neural network including one or more causally masked self-attention layers, e.g., such that at each time step, the self-attention neural network layers only see past inputs in the sequence of processed inputs. The self-attention layers may map queries and a set of key-value pairs derived from each input to the self-attention layer (e.g., all vectors) to an output from which the output of the self-attention layer is derived. The output may be computed as a weighted sum of values ​​weighted by a similarity function of the query for each key.

[0067] If the action selection neural network 102 is a causal transformer, "processing" the input sequence 112 may refer to processing the entire sequence to recalculate the hidden states of previous tokens in the sequence, accessing cached hidden states from memory, and only calculating the hidden state of the last token in the sequence, or using any other techniques to lengthen the context window of attention and reduce inference latency, or both.

[0068] As another example, the action selection neural network 102 may be a recurrent neural network (RNN), i.e., a neural network including one or more recurrent neural network layers. For example, the neural network 102 may be a long short-term memory (LSTM) neural network or a gating recurrent unit (GRU) neural network.

[0069] If the action selection neural network 102 is a recurrent neural network, "processing" the input sequence 112 may refer to processing the entire sequence to recalculate the hidden states of previous tokens in the sequence, or to accessing the most recently updated hidden state from memory and only updating the hidden state by processing the last token in the sequence.

[0070] In some implementations, the environment is a real-world environment, the agent is a mechanical agent that interacts with the real-world environment, such as a robot that operates or navigates in the environment, or an autonomous or semi-autonomous land, air, or sea vehicle, and the action is an action performed by the mechanical agent in the real-world environment to perform a task. For example, the agent may be a robot that interacts with the environment to accomplish a particular task, such as locating an object of interest in the environment, or moving an object of interest to a specified location in the environment, or navigating to a specified location in the environment.

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

[0072] In these implementations, an action can be a control signal for controlling a robot or other mechanical agent, e.g., a torque or higher-level control command for a robot's joint, or a control signal for controlling an autonomous or semi-autonomous land, air, or sea vehicle, e.g., a torque or higher-level control command for a control surface or other control element, e.g., a steering control element of the vehicle. The control signal can include, for example, position, velocity, or force / torque / acceleration data for one or more joints of a robot or part of the other mechanical agent. The control signal can also be, 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, whose control affects the observed state of the environment. For example, in the case of an autonomous or semi-autonomous land, air, or sea vehicle, the control signal can define actions for controlling navigation, e.g., steering, and movement, e.g., braking and / or accelerating the vehicle.

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

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

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

[0076] As one example, the tasks performed by the agents may include tasks to manufacture a product or intermediate products or components thereof. As another example, the tasks performed by the agents may include tasks to control, e.g., minimize, resource usage, such as tasks to control power consumption, or water consumption, or consumption of any materials or consumables used in the manufacturing process.

[0077] Actions may include control actions to control the use of a machine or manufacturing unit to process solid or liquid materials to produce a product, or an intermediate or component thereof, or to control the movement of an intermediate or component of a product within a manufacturing environment, such as between manufacturing units or machines. In general, actions may be any action that affects the observed state of the environment, for example, actions configured to adjust any of the sensed parameters described below. These may include actions to adjust the physical or chemical conditions of a manufacturing unit, or actions to control the movement of a mechanical part of a machine or a joint of a robot. Actions may include actions that impose operating conditions on a manufacturing unit or machine, or actions that result in a change in settings to adjust, control, or turn on or off the operation of a manufacturing unit or machine.

[0078] The reward or return may be related to a metric of the performance of the task. For example, for a task of manufacturing a product, the metric may include the quantity of product produced, the quality of the product, the rate at which the product is produced, or a metric of the physical cost of performing the manufacturing task, such as the amount of energy, material, or other resource used to perform the task. For a task of controlling resource usage, the metric may include any metric of resource usage.

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

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

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

[0082] In general, observations of the state of the environment may include any electronic signals that represent the function of the facility or the equipment at the facility. For example, a representation of the state of the environment may be derived from observations made by any sensors that sense the state of the facility's physical environment, or by any sensors that sense the state of one or more items of equipment, or one or more items of auxiliary control equipment. These include sensors configured to sense electrical conditions such as current, voltage, power, or energy, facility temperature, fluid flow, temperature, or pressure within the facility or within the facility's cooling system, or physical facility configurations such as whether a vent is open or closed.

[0083] The reward or return may be related to a metric of performance of the task. For example, in the case of a task of controlling, e.g., minimizing, resource usage, such as a task of controlling power or water usage, the metric may include any metric of resource usage.

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

[0085] The reward or return may be related to a metric of the performance of the task. For example, in the case of a task of controlling the supply of power to a power distribution grid, the metric may be related to a measure of the power transferred, or a measure of an electrical mismatch between the power supply facility and the grid, such as a voltage, current, frequency, or phase mismatch, or a measure of power or energy losses at the power generation facility. In the case of a task of maximizing the supply of power to a power distribution grid, the metric may be related to a measure of the power or energy transferred to the grid, or a measure of the power or energy losses at the power generation facility.

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

[0087] In another example, the environment may be a chemical synthesis or protein folding environment, where each state is a state of a protein chain or a state of one or more intermediate or precursor chemicals, and the agent is a computer system for determining how to fold a protein chain or synthesize a chemical. In this example, the actions are possible folding actions for folding a protein chain or actions for assembling precursor chemicals / intermediates, and the results to be achieved may include, for example, stabilizing the protein and folding the protein to achieve a specific biological function, or resulting in a valid synthetic pathway for a chemical. In another example, the agent may be a machine agent that indirectly performs or controls protein folding actions or chemical synthesis steps, for example, by automatically controlling synthetic steps selected by the system without human interaction. The observations may include direct or indirect observation of the state of the protein or chemicals / intermediates / precursors and / or may be derived from simulations. Thus, the system may be used to automatically synthesize proteins with specific functions, such as the shape of a binding site, e.g., having a ligand that binds with sufficient affinity for a biological effect that can be used as a drug. For example, it may be a receptor or enzyme agonist or antagonist, or it may be an antibody configured to bind to an antibody target such as a viral coat protein, or a protein expressed on a cancer cell, e.g., to act as an agonist of a particular receptor or to prevent the binding of other ligands and thus the activation of the associated biological pathway.

[0088] Similarly, the environment can be a drug design environment, where each state is a potential drug state and the agent is a computer system for determining drug components and / or drug synthesis pathways. Drugs / synthesis can be designed based on rewards derived from drug targets, for example, in a simulation. As another example, the agent can be a machine agent that performs or controls the synthesis of a drug.

[0089] In some further applications, the environment is a real-world environment and the agent manages the distribution of tasks across computing resources, e.g., on mobile devices and / or in a data center, etc. In these implementations, the action may include assigning a task to a particular computing resource.

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

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

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

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

[0094] The simulated environment may be a simulation of a particular real-world environment and agent. For example, the system may be used to select actions in the simulated environment during training or evaluation of the system, and after training and / or evaluation is complete, and then deployed to control a real-world agent in the particular real-world environment that was the subject of the simulation. This may avoid unnecessary wear and tear on the real-world environment or real-world agent, and may allow for training and evaluation of the control neural network in situations that occur infrequently or are difficult or unsafe to replicate in a real-world environment. For example, the system may be partially trained using a simulation of a machine agent in a simulation of a particular real-world environment, and then deployed to control an actual machine agent in the particular real-world environment. Thus, in such cases, observations of the simulated environment relate to the real-world environment, and actions selected in the simulated environment relate to actions that the machine agent should perform in the real-world environment.

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

[0096] 2 is a flow diagram of an exemplary process 200 for training an action selection neural network. For convenience, process 200 is described as being performed by one or more computer systems located at one or more locations. For example, an action selection system, such as action selection system 100 of FIG. 1, suitably programmed in accordance with this specification, can perform process 200.

[0097] The system obtains a training data set including respective training history sequences for each of a plurality of tasks (step 202).

[0098] The training history sequence for a given task includes sequences of tokens that represent transitions from multiple task episodes executed while training the task's policy through reinforcement learning.

[0099] In other words, each training history sequence contains transitions that were generated at multiple different times during the training of the policy while the task's agents were controlled by the task's policy.

[0100] Generally, the sequence of tokens in the training history sequence includes a respective episode subsequence for each of a plurality of task episodes, and each episode subsequence for each task episode includes a respective transition subsequence for each transition from the task episode.

[0101] Each transition subsequence then includes (i) one or more tokens representing an observation at the transition, (ii) one or more tokens representing an action at the transition, and (iii) one or more tokens representing a reward at the transition.

[0102] So there are a total of N tasks M n When there is a dataset D,

number

number

number

number

number

number

[0103] In general, episode subsequences are ordered within each training history sequence according to the order in which the corresponding task episodes were executed during training.

number

number

[0104] Furthermore, for any given task, the task's policy is represented by a machine learning model, e.g., a neural network or other type of machine learning model, having multiple weights. Thus, because the transition history sequence includes transitions at different times during the training of this machine learning model, and because the values ​​of the weights are repeatedly updated during training, the training history sequence for any given task includes transitions generated while the task's agent is controlled according to multiple different sets of weight values ​​for the multiple weights.

[0105] Thus, any given training history sequence reflects the agent on the task being controlled while using the policy at different stages of training, and can, for example, begin with transitions generated as a result of controlling the agent with a random or near-random policy, and then include subsequent transitions that reflect an improvement in the quality of the actions selected as training progresses.

[0106] In some implementations, the same reinforcement learning algorithm was used to train policies for all tasks. For example, each policy could have been trained using an upper confidence bound (UCB) search method, a policy-based reinforcement learning method such as policy-based actor criticality, an off-policy reinforcement learning method such as off-policy DQN, or any other suitable reinforcement learning method. Thus, the resulting training results in an action selection neural network that "abstracts" or "approximates" the reinforcement learning algorithm. In particular, the training can cause the action selection neural network to implement a more computationally efficient version of the algorithm, i.e., one that requires fewer task episodes to be executed, after training, to achieve or exceed the final performance of the trained policy trained using the algorithm.

[0107] In some other implementations, different reinforcement learning algorithms may have been used to train each policy for different ones of the tasks.

[0108] That is, the training history sequences include both (i) a first training history sequence, the first training history sequence including a sequence of tokens representing transitions from multiple task episodes executed while training a policy for a first task through reinforcement learning using a first reinforcement learning algorithm, and (ii) a second training history sequence of the multiple training history sequences, the second training history sequence including a sequence of tokens representing transitions from multiple task episodes executed while training a policy for a second task through reinforcement learning using a second, different reinforcement learning algorithm. Thus, training of the resulting action selection neural network enables the action selection neural network to generalize across multiple different reinforcement algorithms.

[0109] One exemplary technique for generating a training history sequence is described below with reference to FIG.

[0110] The system then trains the action selection neural network on the training data set by repeatedly performing steps 204 and 206.

[0111] The system selects a subsequence of each training historical sequence from the training data set (step 204).

[0112] In general, the subsequences represent transitions from multiple task episodes represented in the training history sequence, i.e., the subsequences include transition subsequences corresponding to transitions from multiple different episodes of the task represented in the training history sequence.

[0113] For example, the system can randomly sample subsequences of a fixed length, i.e., subsequences containing a fixed number of transition subsequences, with the only constraint that the subsequences contain transition subsequences from multiple episodes. If the fixed length exceeds the maximum number of time steps in an episode, this constraint is automatically satisfied, and the system can randomly sample transition subsequences having the fixed length. The fixed length can be determined, for example, based on the context size of the action selection neural network, where the context size is the number of the largest tokens in any given input sequence to the action selection neural network.

[0114] The system trains the action selection neural network to predict, for each transition in the subsequence, one or more tokens representing the action of the transition conditioned on the tokens preceding the one or more tokens representing the action in the subsequence (step 206).

[0115] Any suitable objective function may be used. By way of example only, the system may train an action selection neural network with a negative log-likelihood loss function L that satisfies:

number

[0116] Thus, as the reinforcement learning policy improves over each history sequence, by accurately predicting actions, the action selection neural network learns to output an improved policy relative to what it sees in the context.

[0117] In some implementations, the system may utilize label smoothing during training, which may include smoothing targets for prediction, e.g., actions at transitions, over other possible targets, e.g., other possible actions.

[0118] In particular, as part of training the negative log-likelihood loss function above, the system can use label smoothing regularization, where the system uses as the target probability distribution for a time step a smoothed distribution that assigns a probability of 1-α to the action at the transition, and α / (k-1) to each other action in the action set, where k is the total number of actions in the set, and α is a positive hyperparameter between 0 and 1.

[0119] 3 is a flow diagram of an exemplary process 300 for generating a training history for a given task. For convenience, process 300 is described as being performed by one or more computer systems located at one or more locations. For example, an action selection system, such as action selection system 100 of FIG. 1, suitably programmed in accordance with this specification, can perform process 300.

[0120] The system trains a policy for a task, e.g., a machine learning model, through reinforcement learning (step 302).

[0121] As part of training the policy, the system repeatedly executes steps 304-310.

[0122] Specifically, the system executes episodes of the task by controlling the task's agents using the task's policy to generate transitions (step 304) and stores the transitions in a replay memory (step 306). The system also samples one or more transitions from the replay memory (step 308) and uses a reinforcement learning algorithm to train a policy on the one or more sampled transitions (step 310).

[0123] In some embodiments, the algorithm is a single-stream algorithm, and steps 304-310 are all performed by the same set of one or more hardware devices.

[0124] In some other implementations, the algorithm is a distributed algorithm, where steps 304 and 306 are performed by multiple actors (“actor computing units”), each of which is implemented on a different set of hardware devices, while steps 308 and 310 are performed by one or more learners (“learner computing units”) implemented on a different set of hardware devices. Generally, multiple actors can perform episodes of a task in parallel with each other during training, i.e., each actor can perform iterations of steps 304 and 306 in parallel with each other actor. Optionally, one or more learners perform steps 308 and 310 asynchronously with the actors performing steps 304 and 306.

[0125] Next, the system generates a training history sequence for the task from the transitions generated as a result of performing the task (step 312).

[0126] If multiple actors perform task episodes when generating the training data, each training history sequence for the task includes episode subsequences of one or more episodes performed by each of the multiple actors.

[0127] In some implementations, the system includes in the training history sequence for a task a respective subsequence for each episode generated during training.

[0128] However, in some other implementations, the system subsamples episodes when generating the training history sequences. Specifically, in each training history sequence for a task, the system includes only the episode subsequence for every kth episode of the task that was performed during training, where k is an integer greater than 1.

[0129] In some cases, rather than subsampling episodes when generating the training history sequence, the system can instead obtain an original training history sequence for a task that includes a respective episode subsequence for each of multiple original task episodes performed during training, and then subsample the original training history sequence by generating a respective historical training sequence that includes only the episode subsequence for the kth original full task episode of the original training history sequence, where k is an integer greater than 1. For example, the original training history could have been generated by another training system.

[0130] In both of these cases, this subsampling can help the action selection neural network become more data-efficient after training when performing reinforcement learning "in context," for example, because the action selection neural network can learn to "extract" more aggressive reinforcement learning algorithms by subsampling.

[0131] FIG. 4 illustrates an example 400 of training the action selection neural network 102.

[0132] Specifically, in example 400, training includes a data generation stage 410 and a model training stage 420.

[0133] In the data generation stage 410, the system generates a training history for each of the N tasks by training a respective policy for each task through reinforcement learning.

number

[0134] As shown in Figure 4 and as described above, the training history for a given task includes data from transitions generated after different amounts of learning progress have occurred during training of the task's policy.

[0135] In the model training stage 420, the system trains the action selection neural network 120 (in the example of FIG. 4, a causal transformer neural network) to predict an action given the above-mentioned episode-wide context. That is, for a transition at time step t, the system selects an action with probability

number

[0136] 5 is a flow diagram of an exemplary process 500 for controlling an agent at a time step using an action selection neural network. For convenience, process 500 is described as being performed by one or more computer systems located at one or more locations. For example, an action selection system, such as action selection system 100 of FIG. 1, suitably programmed in accordance with this specification, can perform process 500.

[0137] The system may execute the process 200 at each time step during a sequence of time steps, for example, at each time step during a task episode.

[0138] The system can also continue to run the process across different task episodes of the same task to perform "in-context" reinforcement learning of the task as described above.

[0139] That is, the system can repeatedly execute process 500 to perform a sequence of task episodes of a task, and during the task episodes, perform "in-context" reinforcement learning about how to perform the task.

[0140] The system receives a current observation characterizing the state of the environment at a time step (step 502).

[0141] The system generates an input sequence of tokens (step 504), which includes (i) one or more tokens representing the current observation, (ii) a respective current transition subsequence for each of one or more current episode transitions, and (iii) a respective previous transition subsequence for each of one or more previous episode transitions.

[0142] Each current episode transition corresponds to a respective previous time step in the current task episode, and each transition subsequence of the current episode transition includes (a) one or more tokens representing observations received in the previous time step, (b) one or more tokens representing actions taken by the agent in response to the observations received in the previous time step, and (c) one or more tokens representing rewards received by the agent in response to taking the actions.

[0143] Each preceding episode transition corresponds to a respective previous time step in a respective preceding task episode performed by the agent before the current task episode, i.e., corresponds to a task episode different from the currently performed task episode, and each preceding transition subsequence of the preceding episode transition includes (a) one or more tokens representing observations received in the previous time step, (b) one or more tokens representing actions performed by the agent in response to the observations received in the previous time step, and (c) one or more tokens representing rewards received by the agent in response to performing the actions.

[0144] Thus, the input sequence includes subsequences for transitions from multiple different episodes, i.e., the current episode and at least one previous episode.

[0145] For example, the system may include subsequences for each transition up to a maximum number of most recently generated transitions in the input sequence, where the maximum number is determined based on the size of the context for the action neural network. If more than the maximum number of transitions have been generated for the current task, the system selects the most recently generated transition represented in the input sequence.

[0146] In some implementations, if fewer than the maximum number of transitions are being generated and the system has access to external data, for example from an expert policy for the task, the system can use the external data to "prompt" the action selection neural network by including in the input sequence subsequences that represent transitions from the expert data in addition to subsequences from the historical data that were generated as a result of controlling the agent using the action selection neural network.

[0147] The system processes the input sequence of tokens using an action selection neural network to generate a policy output for the time step (step 506), and uses the policy output to select an action (step 508).

[0148] The system then causes the agent to perform the selected action (step 510).

[0149] Generally, the agent is controlled by keeping the values ​​of the parameters of the action selection neural network fixed to their trained values, e.g., determined by training the action selection neural network as described above.

[0150] Figure 6 shows an example 600 of attention maps generated by an action selection neural network across time steps when controlling an agent for a new task. In the example of Figure 6, each episode is 50 steps long and the maximum size is a transition from 200 time steps; the left column 610 shows the attention map from steps 0 to 200, and the right column 620 shows the attention map from steps 1700 to 1900. White and gray colors correspond to low and high attention, respectively. From these patterns, it is clear that the action selection neural network attends to the token across several episodes to predict its next action.

[0151] Figure 7 shows example results achieved by the described technique ("AD") against two baselines ("ED" and "Source") on four different tasks. As can be seen from Figure 7, AD's consistent in-context reinforcement effectively learns all four tasks, while being more data-efficient than either of the baselines. Figure 7 also demonstrates the performance upper bound of AD, which is achieved by an online RL algorithm ("RL") that interacts with the environment both during training and action. 2 We show the asymptotic performance of

[0152] In particular, in the example of Figure 7, one of the two baselines is a Source RL ("Source") algorithm used to generate training data for the neural network used in the AD technique. As can be seen from Figure 7, the described technique results in a policy that learns the task more effectively while being significantly more data-efficient than the Source RL algorithm.

[0153] Of the other two baselines ("ED"), one is an expert extraction technique that uses the same action selection neural network as AD, but is trained only on "expert" trajectories rather than learning history. As can be seen from this figure, AD significantly outperforms ED, for example, because AD uses the learned history during training instead of simply the trajectories of high-performing experts.

[0154] The term "configured" is used herein in connection with systems and computer program components. One or more computer systems configured to perform particular operations or actions means that the systems have installed thereon software, firmware, hardware, or a combination thereof that, when running, causes the systems to perform the operations or actions. One or more computer programs configured to perform particular operations or actions means that the one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operations or actions.

[0155] Embodiments of the subject matter and functional operations described herein can be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, such as the structures disclosed herein and their structural equivalents, or one or more combinations thereof. Embodiments of the subject matter described herein can also 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-access or serial-access memory device, or one or more combinations thereof. Alternatively, or in addition, the program instructions can be encoded on an artificially generated transmitted signal, such as 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.

[0156] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or multiple computers. An apparatus may be or further include special-purpose logic circuitry, such as an FPGA (field-programmable gate array) or an ASIC (application-specific integrated circuit). An apparatus optionally includes, 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 one or more combinations thereof).

[0157] A computer program, which may be called or described as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted, declarative or procedural, and can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in part 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, or in multiple cooperating files, e.g., files that store one or more modules, subprograms, or portions of code. A computer program can be deployed to be executed on one computer or to be executed on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communications network.

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

[0159] 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., FPGAs or ASICs), or a combination of special purpose logic circuitry and one or more programmed computers.

[0160] A computer suitable for executing a computer program can be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. Typically, the central processing unit receives instructions and data from a read-only memory or a random-access memory, or both. The basic components of a computer are a central processing unit for 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. Typically, a computer also includes, or is operatively coupled to, one or more mass storage devices for storing data, such as, for example, magnetic, magneto-optical, or optical disks. However, a computer need not include such devices. Furthermore, a computer can be incorporated in other devices, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, such as a universal serial bus (USB) flash drive, to name just a few.

[0161] Computer-readable media suitable for storing computer program instructions and data include, by way of example, all types of non-volatile memory, media, and memory devices, such as 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.

[0162] To interact with a user, embodiments of the subject matter described herein can be implemented in a computer that has a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, and a keyboard and pointing device, such as a mouse or trackball, through which the user can provide input to the computer. Other types of devices can also be used to interact with the user. For example, feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback, and input from the user can be received in any form, including acoustic, speech, or tactile input. Furthermore, a computer may interact with a user by sending documents to a device used by the user 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 types of messages to a personal device (e.g., a smartphone running a messaging application) and then receiving a reply message from the user.

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

[0164] The machine learning model can be implemented and deployed using a machine learning framework, such as the TensorFlow framework or the Jax framework.

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

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

[0167] While the specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what can be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Certain 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 individually or in any suitable subcombination. Furthermore, even if features are described above as functioning in a particular combination and originally claimed as such, one or more features from a claimed combination may, in some cases, be deleted from that combination, and the claimed combination may be directed to a subcombination or variations of the subcombination.

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

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

[0170] Aspects of the present disclosure may be as described in the following clauses.

[0171] Clause 1. A method implemented by one or more computers, said method comprising: obtaining, for each of a plurality of tasks, a training data set including a respective training history sequence; each training history sequence includes a sequence of tokens representing transitions from multiple task episodes executed while training a policy for said task through reinforcement learning; the sequence of tokens includes a respective episode subsequence for each of the plurality of task episodes; The respective episode subsequence for each of the task episodes comprises, for each transition from the task episode: (i) one or more tokens representing observations at said transitions; (ii) one or more tokens representing the actions of the transition; and (iii) one or more tokens representing rewards at said transition; obtaining a transition subsequence including each transition subsequence including training an action selection neural network on the training data set, the training comprising, in each of a plurality of training steps: selecting, from the training data set, a sub-sequence of each training history sequence from the training data set, the sub-sequence representing a transition from a plurality of the task episodes represented in the training history sequence; and training the action selection neural network to predict, for each transition of the subsequence, the one or more tokens representing the action of the transition conditioned on the tokens preceding the one or more tokens representing the action of the subsequence.

[0172] Clause 2. The method of clause 1, wherein the episode subsequences are ordered within each training history sequence according to the order in which the corresponding task episodes were executed during the training.

[0173] Clause 3. The method of any one of clauses 1 to 2, wherein each training history sequence includes transitions generated at multiple different times during the training of the policy while an agent for the task was controlled by the policy for the task.

[0174] Clause 4. The method of clause 3, wherein the policy for the task is represented by a machine learning model having a plurality of weights, and the training history sequence includes transitions generated while the agent for the task was controlled according to a set of a plurality of different weight values ​​for the plurality of weights.

[0175] Clause 5. A first training history sequence of the plurality of training history sequences includes a sequence of tokens representing transitions from a plurality of task episodes executed while training a policy for a first task through reinforcement learning using a first reinforcement learning algorithm; 5. The method of any one of clauses 1 to 4, wherein a second training history sequence of the plurality of training history sequences includes a sequence of tokens representing transitions from a plurality of task episodes executed while training a policy for a second task through reinforcement learning using a second, different reinforcement learning algorithm.

[0176] Clause 6. Obtaining a training data set including a respective training history sequence for each of a plurality of tasks includes, for each of one or more of the tasks: training the policy for the task through reinforcement learning, the training comprising the following operations: executing an episode of the task by controlling an agent for the task using the policy for the task to generate transitions; storing said transitions in a replay memory; sampling one or more transitions from the replay memory; and 6. The method of any one of clauses 1 to 5, comprising repeatedly training the policy on the one or more sampled transitions.

[0177] Clause 7. Obtaining a training data set including a respective training history sequence for each of a plurality of tasks includes, for each of the one or more tasks: The method of clause 6, further comprising: in each training history sequence for the task, including only episode subsequences for every kth episode of the task performed during the training, where k is an integer greater than 1.

[0178] Clause 8. Repeatedly performing the following operations: executing multiple episodes of the task in parallel using multiple actor computing units; 8. The method of claim 6 or claim 7, wherein the respective training history sequences for the tasks include episode subsequences for one or more episodes performed by each of the plurality of actors.

[0179] Clause 9. Obtaining a training data set including a respective training history sequence for each of a plurality of tasks includes, for each of one or more of the tasks: obtaining an original training history sequence for the task, the original training history sequence including a respective episode subsequence for each of a plurality of original task episodes performed during the training; and 9. The method of any one of clauses 1 to 8, comprising generating the respective history training sequences including only episode subsequences for every kth original task episode of the original training history sequence, where k is an integer greater than 1.

[0180] Clause 10. In one or more of said plurality of training steps, training the action selection neural network to predict, for each transition of the subsequence, the one or more tokens representing the action at the transition conditioned on the tokens preceding the one or more tokens representing the action in the subsequence; 10. The method of any one of clauses 1 to 9, comprising training the action selection neural network with label smoothing.

[0181] Clause 11. A method implemented by one or more computers for controlling an agent to perform a sequence of task episodes of a task, the method comprising, for each of a plurality of time steps in a sequence of time steps of a current task episode: receiving a current observation characterizing a state of the environment at the time step; (i) one or more tokens representing said current observation; (ii) a respective current transition subsequence for each of one or more current episode transitions, each current episode transition corresponding to a respective previous time step in the current task episode, the respective transition subsequence for the current episode transition including: (a) one or more tokens representing observations received at the previous time step; (b) one or more tokens representing actions performed by the agent in response to the observations received at the previous time step; and (c) one or more tokens representing rewards received by the agent in response to performing the actions; and (iii) generating an input subsequence of tokens that includes a respective prior transition subsequence for each of one or more prior episode transitions, each prior episode transition corresponding to a respective prior time step in a respective prior task episode performed by the agent prior to the current task episode, and each prior transition subsequence for a prior episode transition including: (a) one or more tokens representing an observation received at the prior time step; (b) one or more tokens representing an action performed by the agent in response to the observation received at the prior time step; and (c) one or more tokens representing a reward received by the agent in response to performing the action; processing the input sequence of tokens using the action selection neural network to generate a policy output for the time step; selecting an action using the policy output; and causing the agent to perform the selected action.

[0182] Clause 12. The method of clause 11, wherein controlling an agent to perform a sequence of task episodes of a task comprises controlling the agent while keeping values ​​of the parameters of the action selection neural network fixed to training values ​​determined by training the action selection neural network by performing each of the actions described in any one of clauses 1 to 10.

[0183] Clause 13. The method of clause 11 or clause 12, wherein the action selection neural network is a causal transformer neural network.

[0184] Clause 14. The method of clause 11 or clause 12, wherein the action selection neural network is a recurrent neural network.

[0185] Clause 15. A method according to any one of clauses 11 to 14, wherein the input sequence comprises the preceding transition subsequence followed by the current transition subsequence followed by the one or more tokens representing the current observation, and the preceding transition subsequence and the current transition subsequence are ordered within the input sequence according to the order in which the corresponding transitions were generated.

[0186] Clause 16. A method according to any one of clauses 1 to 15, wherein each task requires said policy to control a respective agent interacting with a respective environment, and for one or more of said tasks, said respective agent is a machine agent interacting with a real-world environment.

[0187] Clause 17. The method of clause 16, wherein the machine agent is a robot.

[0188] Clause 18. A method according to any one of clauses 1 to 17, wherein each task requires said policy to control a respective agent interacting with a respective environment, and for one or more of said tasks, said environment is a real-world environment of a service facility including a plurality of items of electronic equipment, and said agent is an electronic agent configured to control the operation of said service facility.

[0189] Clause 19. A method according to any one of clauses 1 to 18, wherein each task requires said policy to control a respective agent interacting with a respective environment, and wherein for one or more of said tasks, said environment is a real-world manufacturing environment for producing a product, and said agents include electronic agents configured to control manufacturing units or machines operating to produce said product.

[0190] Clause 20. A 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 the respective method operations of any one of clauses 1 to 19.

[0191] Clause 21. One or more non-transitory computer storage media storing instructions that, when executed by one or more computers, cause said one or more computers to perform the respective method operations of any one of clauses 1 to 19.

Claims

1. 1. A method implemented by one or more computers, the method comprising: obtaining, for each of a plurality of tasks, a training data set including a respective training history sequence; each training history sequence includes a sequence of tokens representing transitions from multiple task episodes executed while training a policy for said task through reinforcement learning; the sequence of tokens includes a respective episode subsequence for each of the plurality of task episodes; The respective episode subsequence for each of the task episodes comprises, for each transition from the task episode: (i) one or more tokens representing observations at said transitions; (ii) one or more tokens representing the action of the transition; and (iii) one or more tokens representing rewards at said transition; obtaining a transition subsequence including each transition subsequence including training an action selection neural network on the training data set, said training comprising: selecting, in each of a plurality of training steps, from the training data set, a subsequence of a respective training history sequence from the training data set, the subsequence representing a transition from a plurality of the task episodes represented in the training history sequence; and training the action selection neural network to predict, for each transition of the subsequence, the one or more tokens representing the action of the transition conditioned on the tokens preceding the one or more tokens representing the action of the subsequence.

2. The method of claim 1 , wherein the episode subsequences are ordered within each training history sequence according to the order in which the corresponding task episodes were performed during the training.

3. 3. The method of claim 1, wherein each training history sequence includes transitions generated at multiple different times during the training of the policy while an agent for the task was controlled by the policy for the task.

4. 4. The method of claim 3 , wherein the policy for the task is represented by a machine learning model having a plurality of weights, and the training history sequence includes transitions generated while the agent for the task was controlled according to a plurality of different sets of weight values ​​for the plurality of weights.

5. a first training history sequence of the plurality of training history sequences including sequences of tokens representing transitions from a plurality of task episodes executed while training a policy for a first task through reinforcement learning using a first reinforcement learning algorithm; 5. The method of claim 1, wherein a second training history sequence of the plurality of training history sequences comprises sequences of tokens representing transitions from a plurality of task episodes executed while training a policy for a second task through reinforcement learning using a second, different reinforcement learning algorithm.

6. Obtaining a training data set including a respective training history sequence for each of a plurality of tasks may include, for each of one or more of the tasks: training the policy for the task through reinforcement learning, the training comprising: executing an episode of the task by controlling an agent for the task using the policy for the task to generate transitions; storing said transitions in a replay memory; sampling one or more transitions from the replay memory; and The method of claim 1 , further comprising: repeatedly performing operations to train the policy on the one or more sampled transitions.

7. Obtaining a training data set including a respective training history sequence for each of a plurality of tasks may include, for each of the one or more tasks:

7. The method of claim 6, further comprising: in the respective training history sequence for the task, including only episode subsequences for every kth episode of the task performed during the training, where k is an integer greater than 1.

8. Repeatedly performing the operations executing multiple episodes of the task in parallel using multiple actor computing units; The method of claim 6 or claim 7, wherein the respective training history sequences for the tasks include episode subsequences for one or more episodes performed by each of the plurality of actors.

9. Obtaining a training data set including a respective training history sequence for each of a plurality of tasks may include, for each of one or more of the tasks: obtaining an original training history sequence for the task, the original training history sequence including a respective episode subsequence for each of a plurality of original task episodes performed during the training; and 9. The method of claim 1, further comprising: generating the respective historical training sequences including only episode subsequences for every kth original task episode of the original training historical sequence, where k is an integer greater than 1.

10. In one or more of the training steps, training the action selection neural network to predict, for each transition of the subsequence, the one or more tokens representing the action at the transition conditioned on the tokens preceding the one or more tokens representing the action in the subsequence; The method of claim 1 , comprising training the action selection neural network with label smoothing.

11. 1. A method implemented by one or more computers for controlling an agent to perform a sequence of task episodes of a task, the method comprising: for each of a plurality of time steps in a sequence of time steps of a current task episode: receiving a current observation characterizing a state of the environment at the time step; generating an input subsequence of tokens, said input subsequence of tokens comprising: (i) one or more tokens representing the current observation; (ii) a respective current transition subsequence for each of one or more current episode transitions, each current episode transition corresponding to a respective previous time step in the current task episode, the respective transition subsequence for the current episode transition including: (a) one or more tokens representing observations received at the previous time step; (b) one or more tokens representing actions performed by the agent in response to the observations received at the previous time step; and (c) one or more tokens representing rewards received by the agent in response to performing the actions; and (iii) a respective prior transition subsequence for each of one or more prior episode transitions, each prior episode transition corresponding to a respective prior time step in a respective prior task episode performed by the agent prior to the current task episode, the respective prior transition subsequence for the prior episode transition including: (a) one or more tokens representing observations received in the prior time step; (b) one or more tokens representing actions performed by the agent in response to the observations received in the prior time step; and (c) one or more tokens representing rewards received by the agent in response to performing the actions; generating, processing the input sequence of tokens using the action selection neural network to generate a policy output for the time step; using the policy output to select an action; and causing the agent to perform the selected action.

12. 12. The method of claim 11, wherein controlling an agent to perform a sequence of task episodes of a task comprises controlling the agent while keeping values ​​of the parameters of the action selection neural network fixed to training values ​​determined by training the action selection neural network by performing the operations of any one of claims 1 to 10.

13. The method of claim 11 or claim 12, wherein the action selection neural network is a causal transformer neural network.

14. The method of claim 11 or claim 12, wherein the action selection neural network is a recurrent neural network.

15. 15. The method of claim 11, wherein the input sequence includes the preceding transition subsequence followed by the current transition subsequence followed by the one or more tokens representing the current observation, the preceding transition subsequence and the current transition subsequence being ordered in the input sequence according to the order in which the corresponding transitions were generated.

16. 16. The method of claim 1, wherein each task requires the policy to control a respective agent that interacts with a respective environment, and wherein for one or more of the tasks, the respective agent is a machine agent that interacts with a real-world environment.

17. The method of claim 16 , wherein the machine agent is a robot.

18. 18. The method of claim 1, wherein each task requires the policy to control a respective agent interacting with a respective environment, and for one or more of the tasks, the environment is a real-world environment of a service facility including a plurality of items of electronic equipment, and the agent is an electronic agent configured to control operation of the service facility.

19. 19. The method of claim 1, wherein each task requires the policy to control a respective agent interacting with a respective environment, and wherein for one or more of the tasks, the environment is a real-world manufacturing environment for producing a product, and the agents include electronic agents configured to control manufacturing units or machines operating to produce the product.

20. A system comprising one or more computers; and 20. A system comprising: 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 the operations of the respective methods described in any one of claims 1 to 19.

21. 20. 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 the respective method operations of any one of claims 1 to 19.

Images