Large scale retrieval enhanced reinforcement learning

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for controlling reinforcement learning agents in an environment to perform tasks. In one aspect, a method includes maintaining a retrieval data set that stores a plurality of historical observations and a respective associated context for each historical observation; receiving a current observation characterizing a current state of the environment; selecting one or more historical observations from the plurality of historical observations; processing, using an encoder neural network, an encoder network input comprising (i) a current observation and (ii) selected one or more historical observations and corresponding associated contexts according to the current value of the encoder network parameter to generate a potential state representation of a current state of the environment; and using the potential state representation to determine an action to be performed by the agent in response to the current observation.

Description

CROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims priority to U.S. Provisional Application No. 63 / 344,008, filed on May 19, 2022. The disclosure of this prior application is considered part of and incorporated by reference into the disclosure of this application. Background Art

[0002] This specification relates to reinforcement learning.

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

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

[0005] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict outputs from received inputs. Some neural networks are deep neural networks that include one or more hidden layers in addition to the output layer. The output of each hidden layer is used as the input to the next layer in the network, the next hidden layer or the output layer. Each layer of the network generates an output from the received inputs based on the current values ​​of the corresponding set of parameters. Summary of the invention

[0006] The specification generally describes a reinforcement learning system that uses a retrieval-augmented action selection process to control an intelligent agent that interacts with an environment. Specifically, a reinforcement learning system as described in the specification enhances the action selection process by enhancing each inference computation of an action selection output based on a current observation of the environment by retrieving relevant empirical data from a retrieval dataset using a learned retrieval process, and then processing data representing the retrieved data to generate an action selection output.

[0007] According to one aspect, a computer-implemented method for controlling a reinforcement learning agent in an environment to perform a task is provided, the method comprising: maintaining a retrieval dataset storing multiple historical observations, the retrieval dataset storing multiple historical observations and corresponding associated contexts for each historical observation, wherein each historical observation represents a corresponding historical state of the environment, and the corresponding associated context for each historical observation represents one or more actions performed in response to the historical observation; receiving a current observation representing a current state of the environment; selecting one or more historical observations from the multiple historical observations, comprising: generating a query vector for the current observation, comprising: processing the current observation using an embedding neural network to generate an embedding, and processing the embedding to generate a query vector, and using the query vector to select one or more historical observations from the multiple historical observations; and using an encoder neural network to process an encoder network input according to current values ​​of encoder network parameters to generate a latent state representation of the current state of the environment, the encoder network input comprising (i) the current observation and (ii) the selected one or more historical observations and their corresponding associated contexts; and using the latent state representation to determine an action to be performed by the agent in response to the current observation.

[0008] The associated context of each historical observation may include one or both of the following: historical actions performed in response to the historical observation, or historical rewards received in response to performing the historical actions.

[0009] The associated context of each historical observation may further include one or both of the following: a historical final observation characterizing a terminal state in a trajectory including the historical observation, a historical final reward received by the agent in the trajectory including the historical observation.

[0010] Processing the embeddings to generate the query vector may include processing the embeddings using principal component analysis.

[0011] Each historical observation stored in the retrieval dataset can be associated with a corresponding key vector generated by using an embedding neural network.

[0012] Selecting one or more historical observations from multiple historical observations may include: selecting one or more key vectors from key vectors associated with multiple historical observations based on: the corresponding distances of the key vectors from a query vector of a current observation; and selecting corresponding historical observations from multiple historical observations associated with the selected one or more key vectors as the selected one or more historical observations.

[0013] Selecting the one or more key vectors may include selecting the one or more key vectors having the smallest Euclidean distance to the currently observed query vector.

[0014] Selecting one or more historical observations from the plurality of historical observations using the query vector may include using an approximate nearest neighbor matching technique.

[0015] Selecting one or more historical observations from the plurality of historical observations using the query vector may include using a vector similarity search technique.

[0016] The vector similarity search technique may include a maximum inner product search technique.

[0017] The maximum inner product search technique may include a maximum inner product search technique based on anisotropic quantization.

[0018] Some or all of the plurality of historical observations stored in the retrieval dataset may be generated as a result of the reinforcement learning agent performing one or more different tasks while interacting with the environment.

[0019] Some or all of the plurality of historical observations stored in the retrieval data set may each have a different dimension, a different modality, or both than the current observation.

[0020] Using the latent state representation to determine an action to be performed by the agent in response to a current observation may include: performing a look-ahead search starting from a current state and searching a latent state representation space guided by a policy network output of a policy neural network, the policy neural network being configured to process, at each time step in a plurality of time steps, a corresponding latent state representation of the time step to generate a corresponding policy network output for the time step.

[0021] The look-ahead search may include a Monte Carlo Tree Search (MCTS) technique.

[0022] The policy network outputs can include: a predicted policy output, which defines a distribution of scores over a set of possible actions; and a predicted value output, which defines a value representing the predicted long-term reward received by the agent due to the environment being in its current state.

[0023] Using the latent state representation to determine an action to be performed by the agent in response to a current observation may include selecting an action to be performed by the agent from a set of possible actions by using a policy network output of a policy neural network, the policy neural network being configured to process the latent state representation of the current state of the environment to generate the policy network output.

[0024] The tasks may include one of the following: an object manipulation task or an environment navigation task.

[0025] The agent may be a mechanical agent, the environment is a real-world environment, and the observations may include data from one or more sensors configured to sense the real-world environment.

[0026] The encoder neural network may include: a first sub-network configured to process a current observation to generate a first encoded representation; a second sub-network configured to process each selected historical observation of the one or more selected historical observations and its corresponding associated context to generate a corresponding second encoded representation; and a third sub-network configured to process the first encoded representation and the corresponding second encoded representation to generate a latent state representation.

[0027] The encoder neural network may have been trained via reinforcement learning on a training dataset that stores multiple trajectories generated as a result of the reinforcement learning agent interacting with the environment.

[0028] The retrieval dataset can be the same as the training dataset.

[0029] The retrieval data set may be different from and larger than the training data set.

[0030] The encoder neural network may have been trained via model-based reinforcement learning.

[0031] According to another aspect, one or more computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the above-described method aspects are provided.

[0032] According to another aspect, a system is provided that includes one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform corresponding operations of the above-described method aspects.

[0033] The subject matter described in this specification can be implemented in specific embodiments to achieve one or more of the following advantages. The technology disclosed in this specification utilizes an information retrieval process to enhance an agent control process using a policy neural network, which information retrieval process additionally provides relevant context information obtained from a retrieval data set. Specifically, by dynamically retrieving relevant information using fast and efficient approximate nearest neighbor matching, the described information retrieval process is easily extended to facilitate large-scale, context-sensitive searches within tens of millions of pieces of information. Therefore, the retrieval process greatly expands the scale of information that the agent control process can access. In principle, this information can be used to train the system, encoding the information in the weights of the system, but it has been found that the retrieval process as described can be used to improve computational and memory efficiency. Some embodiments of the described technology also perform a look-ahead search, and this is combined synergistically with the described retrieval process.

[0034] The retrieval process can be trained end-to-end with a policy neural network, and thus allows the policy neural network to incorporate and utilize relevant information in a learned, i.e., rather than hard-coded, manner when used to generate a policy output for agent control. In addition, new information can be added to the retrieval dataset as it becomes available to continuously enhance the information that can be utilized by the policy neural network, thereby further improving the agent control performance without actually retraining the policy neural network. In some examples, the retrieval process can help control the agent to achieve its task goal faster and more efficiently. In some other examples, the retrieval process can help control the agent to achieve an optimized task goal, i.e., maximizing the expected reward to be received by the agent.

[0035] The details of one or more embodiments of the subject matter of this specification 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 DESCRIPTION OF THE DRAWINGS

[0036] Figure 1 An example reinforcement learning system is shown.

[0037] Figure 2 Is Figure 1 An example illustration of the actions performed by a reinforcement learning system.

[0038] Figure 3 is a flowchart of an example process for training an encoder neural network and a policy neural network.

[0039] Figure 4 is a flow chart of an example process for controlling a reinforcement learning agent.

[0040] Figure 5 yes Figure 4 A flowchart of the sub-steps of one of the steps of a process.

[0041] Figure 6 Quantitative examples of performance gains that can be achieved by an agent controlled using the reinforcement learning system described in this specification are shown.

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

[0043] This specification describes a reinforcement learning system that controls an agent to interact with an environment by processing, at each of a plurality of time steps, data representing the current state of the environment at that time step (i.e., an "observation") to select an action to be performed by the agent.

[0044] Once the reinforcement learning system selects an action to be performed by the agent, the reinforcement learning system can cause the agent to perform the selected action. For example, the system can instruct the agent and the agent can perform the selected action. As another example, the system can directly generate control signals for one or more controllable elements of the agent. As yet another example, the system can transmit data specifying the selected action to a control system of the agent, which controls the agent to perform the action. Typically, the agent performing the selected action causes the environment to transition to a different state.

[0045] The systems described herein are broadly applicable and are not limited to one particular implementation. However, for illustrative purposes, a few example implementations are described below.

[0046] In some embodiments, the environment is a real-world environment, the agent is a mechanical agent that interacts with the real-world environment, such as a robot or an autonomous or semi-autonomous land, air, or sea vehicle that operates in or navigates through the environment, and the action is an action taken by the mechanical agent in the real-world environment to perform a task. For example, the agent can be a robot that interacts with the environment to complete a specific task, such as to locate an object of interest in the environment or to move an object of interest to a specified location in the environment or to navigate to a specified destination in the environment.

[0047] In these embodiments, observations may include, for example, one or more of the following: images, object position data, and sensor data for capturing observations when the agent interacts with the environment, such as sensor data from images, distance or position sensors, or from actuators. For example, in the case of a robot, observations may include data characterizing the current state of the robot, such as one or more of the following: joint positions, joint velocities, joint forces, torques, or accelerations, such as torque feedback for gravity compensation, and the global or relative posture of an item held by the robot. In the case of a robot or other mechanical agent or vehicle, observations may similarly include one or more of the position, linear or angular velocity, force, torque, or acceleration, and global or relative posture of one or more components of the agent. Observations may be defined in 1, 2, or 3 dimensions, and may be absolute and / or relative observations. Observations may also include, for example, sensed electronic signals, such as motor currents or temperature signals; and / or, for example, image or video data from a camera or LIDAR sensor, such as data from a sensor of the agent or data from a sensor located separately from the agent in the environment.

[0048] In these embodiments, the action can be a control signal for controlling a robot or other mechanical agent, for example, a torque or higher-level control command to a joint of the robot, or an autonomous or semi-autonomous land, air, or marine vehicle, for example, a torque to a control surface or other control element, for example, a steering control element of a vehicle, or a higher-level control command. The control signal may include, for example, position, velocity, or force / torque / acceleration data of one or more joints of a robot or a component of another mechanical agent. The control signal may also or instead include electronic control data, such as motor control data, or more generally data for controlling one or more electronic devices within an environment, the control of which has an effect on the observed state of the environment. For example, in the case of an autonomous or semi-autonomous land, air, or marine vehicle, the control signal may define actions that control navigation, such as steering, and movement, such as braking and / or acceleration of the vehicle.

[0049] In some embodiments, 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 can be a simulation of a robot or a vehicle, and the reinforcement learning system can be trained on the simulation and then, once trained, used in the real world.

[0050] In some embodiments, the environment is a real-world manufacturing environment for manufacturing products or foods such as chemical, biological or mechanical products. As used herein, "manufacturing" products also includes refining starting materials to create products, or processing starting materials, for example, to remove pollutants, to produce clean or recycled products. A manufacturing plant may include multiple manufacturing units, such as containers for chemical or biological substances, or machines for processing solids or other materials, such as robots. The manufacturing units are configured so that intermediate versions or components of the product can be moved between manufacturing units during the manufacture of the product, such as via pipelines or mechanical transport. As used herein, the manufacture of products also includes manufacturing food products by kitchen robots.

[0051] The agent may include an electronic agent configured to control a manufacturing cell or a machine such as a robot that operates to manufacture a 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 cells or machines, or to control the movement of intermediate versions or components of a product between manufacturing cells or machines.

[0052] As one example, the tasks performed by the agent may include the task of manufacturing a product or an intermediate version or component thereof. As another example, the tasks performed by the agent may include the task of controlling - e.g., minimizing - the use of a resource, such as the task of controlling the consumption of electricity or water or the consumption of any material or consumable used in the manufacturing process.

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

[0054] The reward or reward may relate to a measure of the performance of the task. For example, in the case of a task to manufacture a product, the measure may include a measure of the quantity of the product manufactured, the quality of the product, the speed at which the product is produced, or the physical cost of performing the manufacturing task, such as a measure of the amount of energy, material, or other resource used to perform the task. In the case of a task to control the use of resources, the matrix may include any measure of the use of the resource.

[0055] In general, observations of environmental states may include any electronic signals representing the operation of electronic and / or mechanical items of equipment. For example, representations of environmental states may be derived from observations made by sensors that sense the state of a 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 may be configured to sense mechanical movement or force, pressure, temperature; electrical conditions such as current, voltage, frequency, impedance; the amount, level, flow / movement rate, or flow / movement path of one or more materials; physical or chemical conditions such as physical state, shape, or configuration or chemical state such as pH; configuration of a unit or machine, such as a mechanical configuration or valve configuration of a unit or machine; image or video sensors for capturing images or video observations of a manufacturing unit or machine or movement; or any other appropriate 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 posture of one or more components of the machine, such as data characterizing the current state of the machine or robot or an item 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, such as from a camera or LIDAR sensor. Sensors such as these may be part of or located separately from the agents in the environment.

[0056] In some embodiments, the environment is a real-world environment of a service facility including multiple items of electronic equipment, such as a server farm or data center, such as a telecommunications data center, or a computer data center for storing or processing data, or any service facility. The service facility may also include auxiliary control equipment that controls the operating environment of the item of equipment, such as environmental control equipment, such as temperature control, such as cooling equipment, or airflow control or air conditioning equipment. Tasks may include tasks that control - for example, minimize - the use of resources, such as tasks that control power consumption or water consumption. Agents may include electronic agents that are configured to control the operation of an item of equipment, or to control the operation of auxiliary - for example, environmental - control equipment.

[0057] In general, an action may be any action that has an effect on the observed state of the environment, such as an action configured to adjust any of the sensed parameters described below. These may include actions that control an item of equipment or auxiliary control equipment or impose operating conditions on an item of equipment or auxiliary control equipment, such as actions that cause a setting to be changed to adjust, control, or switch on or off the operation of an item of equipment or an item of auxiliary control equipment.

[0058] Typically, observations of the state of the environment may include any electronic signal that represents the operation of the facility or equipment in the facility. For example, representations of the state of the environment may be derived from observations made by any sensor that senses the state of the physical environment of the facility or by any sensor that senses the state of one or more items of equipment or assists in controlling the state of one or more items of equipment. These sensors include sensors configured to sense electrical conditions such as current, voltage, power, or energy; the temperature of the facility; fluid flow, temperature, or pressure within the facility or within the cooling system of the facility; or physical facility configurations, such as whether a vent is open.

[0059] The reward or reward may relate to a measure of the performance of the task. For example, in the case of a task to control - e.g. minimise - the use of a resource - such as a task to control the use of electricity or water - the measure may comprise any measure of the use of the resource.

[0060] In some embodiments, the environment is a real-world environment of a power generation facility, such as a renewable power generation facility, such as a solar farm or a wind farm. The task may include a control task to control the power generated by the facility, for example, to control the transmission of power to a distribution network, for example, to meet demand or reduce the risk of mismatch between elements of the power grid, or to maximize the power generated by the facility. The agent may include an electronic agent configured to control the generation of power by the facility or to couple the generated power to the power grid. These actions may include actions to control the electrical or mechanical configuration of the generator, such as the electrical or mechanical configuration of one or more renewable power generation elements, for example, for controlling the configuration of a wind turbine or one or more solar or mirrors, or the electrical or mechanical configuration of a rotating generator. Mechanical control actions may, for example, include actions to control the conversion of energy input to electrical energy output, such as the efficiency of the conversion or the degree of coupling of energy input to electrical energy output. Electrical control actions may, for example, include actions to control one or more of the voltage, current, frequency, or phase of the generated power.

[0061] The reward or reward may involve a measure of the performance of the task. For example, in the case of a task to control the delivery of electricity to a distribution network, the measure may involve a measure of the power transmitted, or a measure of an electrical mismatch such as a voltage, current, frequency or phase mismatch between the generating facility and the grid, or a measure of power or energy lost in the generating facility. In the case of a task to maximize the delivery of electricity to a distribution network, the measure may involve a measure of power or energy delivered to the grid, or a measure of power or energy lost in the generating facility.

[0062] Typically, observations of environmental conditions may include any electronic signals representing the electrical or mechanical operation of power generation equipment in a power generation facility. For example, representations of environmental conditions may be derived from observations made by any sensor that senses the physical or electrical state of equipment in a power generation facility that is generating electricity, or the physical environment of such equipment, or the condition of auxiliary equipment that supports the power generation equipment. Such sensors may include sensors configured to sense electrical conditions such as current, voltage, power, or energy of the equipment; temperature or cooling of the physical environment; fluid flow; or the physical configuration of the equipment; and observations of electrical conditions of the power grid, such as from local or remote sensors. Observations of environmental conditions may also include one or more predictions about future conditions of the operation of the power generation equipment, such as predictions of future wind levels or solar irradiance or predictions of future electrical conditions of the power grid.

[0063] As another example, the environment can be a chemical synthesis or protein folding environment, such that each state is a corresponding state of a protein chain or one or more intermediates 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, these actions are possible folding actions for folding a protein chain or actions for assembling precursor chemicals / intermediates, and the results to be achieved can include, for example, folding a protein so that the protein is stable and enables it to achieve a specific biological function or provide an efficient synthetic pathway for a chemical. As another example, the agent can be a mechanical agent that performs or controls a protein folding action or chemical synthesis step automatically selected by the system without human interaction. Observations can include direct or indirect observations of the state of a protein or chemical / intermediate / precursor and / or can be obtained from a simulation.

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

[0065] In some further applications, the environment is a real-world environment, and the agent manages the distribution of tasks across computing resources, such as on mobile devices and / or in data centers. In these embodiments, an action may include assigning a task to a specific computing resource.

[0066] As further examples, an action may include presenting an advertisement, an observation may include an ad impression or a click-through count or rate, and a reward may characterize a previous selection of an item or content by one or more users.

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

[0068] As another example, the environment may be an electrical, mechanical, or electromechanical design environment, such as an environment in which the design of an electrical, mechanical, or electromechanical entity is simulated. A simulation environment may be a simulation of a real-world environment in which the entity is intended to work. The task may be to design an entity. Observations may include observations that characterize the entity, i.e., observations of the mechanical shape or electrical, mechanical, or electromechanical configuration of the entity, or observations of parameters or properties of the entity. Actions may include actions that modify the entity, such as actions that modify one or more of the observations. A reward or reward may include one or more measures of the performance of the design of the entity. For example, the reward or reward may relate to one or more physical properties of the entity, such as weight or strength, or to one or more electrical properties of the entity, such as a measure of the efficiency of performing a specific function for which the entity is designed. The design process may include, for example, outputting a design for manufacturing in the form of computer executable instructions for manufacturing the entity. The process may include making the entity according to the design. Thus, the design of the entity may be optimized, for example, by reinforcement learning, and then the optimized design output for manufacturing the entity may be optimized, for example, as computer executable instructions; then the entity with the optimized design may be manufactured.

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

[0070] The simulation environment can be a simulation of a specific real-world environment and agent. For example, the system can be used to select actions in the simulation environment during training or evaluation of the system, and after the training or evaluation or both are completed, it can be deployed to control the real-world agent in the specific real-world environment as the simulation subject. This can avoid unnecessary wear and damage to the real-world environment or the real-world agent, and can allow the training and evaluation of control neural networks in situations that rarely occur or are difficult to recreate or cannot be safely recreated in the real-world environment. For example, the system can be partially trained using a simulation of a mechanical agent in a simulation of a specific real-world environment, and then the system can be deployed to control a real mechanical agent in a specific real-world environment. Therefore, in such a case, the observation of the simulation environment is related to the real-world environment, and the selected actions in the simulation environment are related to the actions to be performed by the mechanical agent in the real-world environment.

[0071] As another example, in some embodiments, the agent includes a digital assistant, such as a smart speaker, smart display, or other device, and the action performed by the agent is an output generated by the digital assistant in response to input from a human user specifying a task to be performed. The output can be provided using natural language, for example, on a display and / or using a speech synthesis subsystem of the digital assistant. The digital assistant can be used, for example, to capture visual - e.g., video - and / or audio observations of the user.

[0072] Optionally, in any of the above embodiments, 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.

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

[0074] Reinforcement learning system 100 controls an agent 102 interacting with environment 104 by selecting an action 106 to be performed by agent 102 and then causing agent 102 to perform the selected action 106, such as by transmitting control data to agent 102 instructing agent 102 to perform action 102. In some other cases, RL system 100 is remote from the agent and transmits control data or other instructions that cause the agent to perform an action.

[0075] Execution of the selected actions 106 by the agent 102 typically causes the environment 104 to transition to successive new states. By repeatedly causing the agent 102 to act in the environment 104, the system 100 can control the agent 102 to complete a specified task.

[0076] The reinforcement learning system 100 includes an encoder neural network 120 and a policy neural network 130, which are used together to control the agent 102. Optionally, in some embodiments, the reinforcement learning system 100 includes or has access to an environment dynamics model that works in conjunction with the encoder neural network 120 and the policy neural network 130 to control the agent 102.

[0077] The reinforcement learning system 100 also includes a retrieval dataset 150 that stores experience data generated as a result of the interaction of the agent 102 or another agent with the environment 104 or with another instance of the environment. Additionally or alternatively, the retrieval dataset 150 stores experience data derived from environment interaction information obtained from any of a variety of other relevant sources, such as a driving log in the case where the agent is an autonomous or semi-autonomous vehicle. For example, the agent may be controlled by the present reinforcement learning system 100 or another control system, including by a human-operated or another machine learning-based control system, by a hard-coded policy that selects actions to be performed by the agent according to pre-programmed logic, or simply by a random policy that selects actions with uniform randomness. Collectively, the experience data stored in the retrieval dataset 150 represents past experience information that controlled the agent to perform the same or different tasks.

[0078] More specifically, the experience data includes a plurality of historical observations obtained during past interactions of the agent 102 or another agent with the environment 104 or with another instance of the environment, and the retrieval data set 150 stores these historical observations in association with corresponding key vectors. Each historical observation characterizes a corresponding historical state of the environment 104 or another instance of the environment. Each historical observation may, but need not, have the same dimensions or the same data modality as the current observation 108. In general, observations with different data modalities may refer to observations based on different types of data representing the state of the environment, for example, based on audio versus (vs) image versus text data.

[0079] For each historical observation, the experience data also includes an associated context, which may include, among other information, any information related to the specified task that the agent 102 was performing in the past, but which may still not be obvious from the historical observation itself. For example, the associated context of each historical observation may include data specifying one or more subsequent historical actions performed in response to the historical observation, one or more subsequent historical rewards received in response to the subsequent historical actions being performed, or both. As another example, the associated context of each historical observation may include a historical final observation that characterizes the terminal state in the trajectory that includes the historical observation, a historical final reward received by the agent in the trajectory that includes the historical observation, or both. As yet another example, the associated context of each historical observation may include all or some subset of the example information mentioned above.

[0080] A trajectory represents a sequence of interactions of an agent with an environment over a sequence of time steps. For example, a trajectory may represent an episode of a specified task over a sequence of time steps during which the agent attempts to perform the specified task. For example, a task episode may continue for a predetermined number of time steps, or until a reward is received indicating that the task has been successfully completed.

[0081] At each of the plurality of time steps, the reinforcement learning system 100 performs a learned information retrieval process to retrieve relevant experience data 152 from the retrieval dataset 150, and provides the retrieved data 152 to the encoder neural network 120, i.e., appended to the current observation 108. As used in this specification, the term "learned" means that a process or value has been adjusted during the training of the system.

[0082] To perform the information retrieval process, the reinforcement learning system 100 uses the embedding neural network 110 to generate a query vector 112 for the current observation 108, and then uses the query vector 112 to select one or more historical observations and their associated contexts from the retrieval dataset 150 as relevant experience data 152 to provide to the encoder neural network 150.

[0083] Specifically, in some embodiments, the selection process utilizes fast (approximate) nearest neighbor matching between query vectors and key vectors to facilitate large-scale retrieval, i.e., selecting relevant experience data from very large retrieval datasets, e.g., datasets including tens of millions or more historical observations and their associated contexts. Thus, the relevant experience data 152 selected by the information retrieval process, by virtue of its diversity, enables better agent control by flexibly associating past experience and relevant contextual information to the current task.

[0084] The encoder neural network 120 is configured to process an encoder network input at each of a plurality of time steps, the encoder network input comprising (i) a current observation 108 characterizing a current state of the environment 104 at a given time step and (ii) relevant experience data 152 to generate a latent state representation 122 of the current state of the environment 104.

[0085] This latent state representation 122 and, in some embodiments, action data (or data derived from the action data, or both) defining each action from the set of possible actions can then be provided as input to a policy neural network 130 for processing to generate a policy network output.

[0086] The policy network output includes a policy output 132 and optionally a value output 133 .

[0087] Value output 133, when provided, represents an estimate of the value that environment 104, in its current state, will successfully perform a specified task. For example, the value estimate may be an estimate of the reward of the task resulting from environment 104 being in its current state as represented by observation 108. Reward refers to a cumulative measure of the "reward" received by the agent, such as a time-discounted sum of the rewards. Value output 133 may be used, for example, to perform a lookahead search using latent state representation 122 to determine an action to be performed by the agent.

[0088] The reinforcement learning system 100 then uses the policy output 132 to select an action 106 to be performed by the agent 102 at the current time step.

[0089] In some implementations, the policy output 132 generated by the policy neural network 130 for a time step is used directly to select an action 106 to be performed by the agent 102 at the current time step. Several examples of using the policy output 132 to select an action 106 to be performed by the agent 102 are described next.

[0090] In one example, the policy output 132 may include a corresponding numeric probability value for each action in a set of possible actions that may be performed by the agent. The system may select actions to be performed by the agent, for example, by sampling actions according to their probability values, or by selecting the action with the highest probability value.

[0091] In another example, the policy output 132 may directly define the actions to be performed by the agent, for example, by defining torque values ​​that should be applied to the joints of the robotic agent.

[0092] In another example, the policy output 132 may include a corresponding Q-value for each action in a set of possible actions that can be performed by the agent. The system may process the Q-values ​​(e.g., using a soft-max function) to generate a corresponding probability value for each possible action, which may be used to select an action to be performed by the agent (as described earlier). The system may also select the action with the highest Q-value as the action to be performed by the agent.

[0093] The Q-value of an action is an estimate of the “reward” that would result from the agent performing the action in response to the current observation and thereafter choosing future actions to be performed by the agent based on the current values ​​of the policy neural network parameters.

[0094] Return refers to a cumulative measure of the "rewards" received by an agent, such as a time-discounted sum of rewards. The agent may receive a corresponding reward at each time step, where the reward is specified by a scalar value and represents, for example, the agent's progress in completing an assigned task.

[0095] In some cases, the system also receives an extrinsic reward for the task from the environment ("task reward"). Typically, the extrinsic reward is a scalar value and characterizes the agent's progress toward completing the task. As a specific example, the extrinsic reward can be a sparse binary reward that is zero unless the task is successfully completed, and is one if the task is successfully completed as a result of the action performed. As another specific example, the extrinsic reward can be a dense reward that measures the agent's progress toward completing the task at the individual observations received during the episodes in which it attempts to perform the task, i.e., such that a non-zero reward can and frequently is received before the task is successfully completed.

[0096] In some cases, the system 100 may select actions to be performed by the agent based on an exploration strategy. For example, an exploration strategy may be - Greedy exploration strategy, where the system chooses a target with probability 1- The policy output 132 selects an action to be performed by the agent and randomly selects an action with probability In this example, is a scalar value between 0 and 1.

[0097] In some other embodiments, rather than using the policy output 132 directly to control the agent 102, the system 100 instead performs a plurality of planning iterations guided by the policy output 132 and, in the case of providing a value output 133 of the policy neural network 130, then selects an action 106 to be performed in response to the current observation 108 based on the results of the planning iterations. Each planning iteration typically involves performing a lookahead search, such as a tree-based lookahead search algorithm, such as a Monte Carlo Tree Search (MCTS) algorithm, to repeatedly (i.e., at each of a plurality of planning steps in each planning iteration) select a sequence of actions that advances the environment from the state represented by the current observation to a new state. This allows the system to search the space of potential state representations starting from the current state represented by the current observation 108. For example, a look-ahead search may be guided by the policy network output by searching a space of latent state representations representing the state of the environment, using the policy output 132 and, when provided, the value output 133 of the policy neural network 130, using the policy output 132 to identify potential future actions, and optionally using the value output 133 to evaluate the state of the environment, such as a terminal search state of the environment.

[0098] The system 100 can use the policy output 132 and, in some cases, the value output 133 together with the environment dynamics model to run a look-ahead search algorithm to determine which actions can be selected to be performed by the agent. The environment dynamics model is a model that, given information at a given time step, is able to make predictions about at least one future time. When used, the environment dynamics model can make predictions about multiple time steps after the given time step. For example, the environment dynamics model can be configured to receive as input a latent state representation 122 corresponding to a current observation 108 and data specifying a current action 106 selected using the policy output 132 and the value output 133, and generate as output a predicted feature latent state representation that will be generated by the encoder neural network 120 for the next observation representing the next state to which the environment 104 transitions.

[0099] The use of lookahead search to control an agent may be referred to as the "MuZero" technique and is described in more detail in Julian Schrittwieser, et al. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model. Nature, 588(7839):604-609, 2020, the entire contents of which are hereby incorporated herein in their entirety. Reference may be made in particular to the MuZero algorithm described in Section 3 (also in Schrittwieser, et al. arXiv:1911.08265).

[0100] Figure 2 Is Figure 1 An example illustration of operations performed by the reinforcement learning system 100 of FIG. 100. In general, the system may respond to receiving a representation corresponding to the current time step by These operations are performed based on the current observation 108 of the current state of the corresponding environment 104.

[0101] The system processes the current observation 108 using an embedding neural network 110 to generate an embedding, and then uses the embedding to generate a query vector 112. The embedding neural network 110 can be implemented with any suitable neural network architecture. Specifically, the embedding neural network 110 can be a feed-forward neural network that includes any suitable number (e.g., 5 layers, 10 layers, or 100 layers) of any suitable neural network layers (e.g., fully connected layers, convolutional layers, attention layers, etc.) and arranged in any suitable configuration (e.g., as a linear sequence of layers with or without residual connections). The embedding of an observation can refer to a representation of an observation as an ordered collection of numerical values, such as a vector or matrix of numerical values.

[0102] The embedding may be an activation output of an intermediate layer in the embedding neural network 110, or a combination of activation outputs of two or more layers in the embedding neural network 110. In some cases, the embedding is used directly as the query vector 112, while in other cases, the embedding is further processed to generate the query vector 112. For example, the system 100 may apply a linear projection to the embedding to reduce the dimensionality of the embedding using principal component analysis or another dimensionality reduction technique (e.g., non-negative matrix factorization, linear discriminant analysis, etc.). The projected embedding may then be used as the query vector 112.

[0103] As mentioned above, the retrieval data set 150 is associated with a corresponding key vector 113a-n to store each historical observation in the plurality of historical observations 151a-n. The embedding neural network 110 is also used to generate these key vectors 113a-n. Similar to how the query vector 112 is generated, the system 100 uses the embedding neural network 110 to pre-process (i.e., before controlling the intelligent agent 102 to perform any task) each historical observation to generate an embedding, and then uses the embedding to generate the corresponding key vector 113.

[0104] The system 100 selects one or more historical observations from the retrieval data set 150 by using the query vector 112 of the current observation 108 and the corresponding key vectors 113a-n of the plurality of historical observations 151a-n. Specifically, this may include calculating the corresponding distance between the corresponding key vector of the historical observation and the query vector 112 of the current observation 108 for each historical observation stored in the retrieval data set 150, and then identifying one or more key vectors based on the corresponding distance, such as identifying the (approximate) nearest neighbor. Thus, one or more historical observations associated with the identified key vector may be selected. For example, the system may select one or more key vectors having the minimum distance to the query vector 112 or satisfying a distance threshold, and the distance may be a Euclidean distance, a Hamming distance, or other types of distances. In this way, the system 100 may select a relatively small number of historical observations, such as four, or another number less than ten, that are most relevant to the current observation 108 from the retrieval data set 150 as relevant historical observations.

[0105] The system 100 provides the current observation 108 and the selected one or more historical observations along with their associated contexts as inputs to the encoder neural network 120, and then processes the inputs using the encoder neural network 120 to generate a latent space representation 122 of the current state of the environment. As just one example, the encoder neural network 120 may generate embeddings of the selected one or more historical observations and their associated contexts, and combine them in a permutation-invariant manner, such as by summing them, and then may concatenate the combination with the embedding of the current observation 108.

[0106] Similar to the embedding neural network 110, the encoder neural network 120 may be implemented with any suitable neural network architecture that enables it to map inputs to a latent space representation. Figure 2In the example of , the encoder neural network 120 includes a first sub-network 120A, a second sub-network 120B, and a third sub-network 120C. A sub-network of a neural network refers to a set of one or more neural network layers in a neural network. Each sub-network may include, for example, one or more fully connected layers, one or more convolutional layers, one or more attention layers, one or more residual blocks, etc. The techniques described herein are not limited to the use of any particular neural network architecture. For example, observations including images can be processed by fully connected layers, convolutional layers, layers, and / or using residual blocks.

[0107] As shown, the first sub-network 120A processes the current observation 108 to generate a first encoded representation. The second sub-network 120B processes each of the four selected historical observations and their respective associated contexts to generate a corresponding second encoded representation. The third sub-network 120C processes a combination of the first encoded representation and the corresponding second encoded representation, such as a sum or average, to generate a latent state representation 122.

[0108] System 100 represents the potential state 122 is provided as input to the policy neural network 130, which is then used to process the input to generate a policy output , and optionally generate value output , the value output It can then be used to select The action to be performed by the agent.

[0109] exist Figure 2 In the example of , the policy neural network 130 is a recurrent neural network that includes one or more recurrent layers, such as a long short-term memory (LSTM) layer or a gated recurrent unit (GRU) layer, followed by a policy head that provides a policy output and a value head that generates a value output. The policy head and the value head can each be a corresponding feedforward neural network.

[0110] The policy neural network 130 at the current time step The latent state representation 122 is processed at the time to update the current hidden state of the policy neural network 130, that is, the current hidden state of the policy neural network 130 that has been generated by processing the previous input state representation at one or more previous time steps is modified, and the updated current hidden state is then processed by the policy head to provide a policy output, and the updated current hidden state is processed by the value head to generate a value output.

[0111] Before using the neural network to control the agent 102 to perform a specified task, the reinforcement learning system 100 or another training system trains the embedding neural network 110, the encoder neural network 120, and the policy neural network 130 on the training data set of the specified task, that is, so that the embedding neural network 110 can be used to generate a query vector 112, the query vector 112 enables the selection of highly relevant and informative relevant experience data 152 from the retrieval data set 150, and the encoder neural network 120 and the policy neural network 130 can effectively use the relevant experience data 152 to enable the selection of actions that can lead to the successful completion of the specified task. In some cases, the training data set is the same as the retrieval data set 150, while in other cases, the training data set is different from the retrieval data set 150, for example, including fewer observations.

[0112] More generally, the training data set stores a plurality of transitions generated as a result of the agent 102 (or another agent) interacting with the environment 104 (or another instance of the environment). Each transition includes (1) a current observation representing the current state of the environment at a time; and optionally, one or more of: (2) a current action performed by the agent in response to the current observation; (3) a next observation representing the next state of the environment after the agent performs the current action, i.e., the state to which the environment transitions as a result of the agent performing the current action; or (4) a reward received in response to the agent performing the current action. As an example, each transition may include a tuple of (1)-(3) or (1)-(4) as data elements.

[0113] Figure 3 is a flow chart of an example process 300 for training an encoder neural network and a policy neural network. For convenience, process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a suitably programmed reinforcement learning system—e.g. Figure 1 The reinforcement learning system 100 of ——can perform process 300.

[0114] The system can continue to perform iterations of process 300 to train the encoder neural network and the policy neural network using different batches of transformations sampled from the training dataset at different iterations of process 300 until a termination criterion for the training of the neural networks has been met, e.g., until parameters have converged, until a threshold amount of wall clock time has elapsed, or until a threshold number of iterations of process 300 have been performed. In some cases, the system distributes training across multiple computing units.

[0115] For each transition included in the batch, the system processes the encoder network input including the current observation included in the transition and determines a reinforcement learning (RL) loss associated with the batch of transitions (step 302). The encoder network input may also include one or more selected historical observations and their corresponding associated contexts, which are selected from a retrieval dataset using a nearest neighbor retrieval process based on the current observation included in the batch.

[0116] The RL loss can generally be determined using any of a variety of offline RL techniques. For example, the system can use the offline actor-critic technique described in Sergey Levine, et al. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020, where the policy neural network is configured as an actor and / or critic. As another example, the system can use the value iteration technique described in Martin Riedmiller. Neuralfitted q iteration–first experiences with a data efficient neural reinforcement learning method. In European conference on machine learning, pages 317–328. Springer, 2005, where the policy neural network is configured as a Q neural network. As yet another example, the system can use a model-based technique, such as employed by the MuZero technique. Model-based reinforcement learning techniques may be techniques in which a model, such as a neural network, predicts future states of an environment, which may then be used to select actions rather than predicting actions directly.

[0117] The system may also apply any of a variety of regularization techniques during RL training of the system to improve network robustness to irrelevant or poor quality historical observations that may be selected. For example, the system may randomly zero out a subset of retrieved historical observations during training. As another example, the system may randomly replace a subset of retrieved historical observations with historical observations selected for different observations. As yet another example, the system may add a regularization loss to the RL loss, such as a mean squared error loss, which is used to regularize the embeddings generated for the retrieved historical observations towards the embedding generated for the current observation.

[0118] The system determines one or more updates to the parameter values ​​of the policy neural network based on the gradients of the RL loss calculated with respect to the parameters of the policy neural network to train the policy neural network to optimize the RL loss (step 304). The system can determine the updates by applying an optimizer to the gradients suitable for training the policy neural network—e.g., an Adam optimizer, an rmsProp optimizer, or a stochastic gradient descent optimizer.

[0119] The system also determines one or more updates to parameter values ​​of the encoder neural network by back-propagating the gradient of the RL loss into the encoder neural network (step 306).

[0120] Because the nearest neighbor retrieval process is non-differentiable, which means that the query and key vector mapping cannot be directly trained end-to-end, instead of jointly learning parameter values ​​for the embedding neural network during RL training, the system 100 can instead pre-train the embedding neural network 110 on the training dataset, e.g., as a standalone neural network or as part of another larger neural network, and then use the pre-trained embedding neural network 110 to generate queries and keys, which are then used in process 300 to perform RL training of the encoder neural network 120 and the policy neural network 130. For example, the system 100 can train an action selection policy model on the transformations stored in the training dataset, e.g., using MuZero techniques (i.e., in the case of the retrieval process), and then use (a portion of) the trained action selection policy model as part of the embedded neural network 110, e.g., between the input layer and an intermediate layer before the policy output layer of the model.

[0121] Figure 4 is a flow chart of an example process 400 for controlling a reinforcement learning agent in an environment to perform a task. For convenience, process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a suitably programmed reinforcement learning system—e.g. Figure 1 The reinforcement learning system 100 of can perform process 400.

[0122] The system maintains a retrieval data set that stores a plurality of historical observations and a corresponding associated context and a corresponding key vector for each historical observation (step 402). The key vector is generated based on the historical observations by using an embedding neural network. The retrieval data set can be generated based on experience data generated as a result of the agent or another agent performing the same or different task in the same or different environment.

[0123] Each historical observation characterizes a corresponding historical state of the environment. As described above, the corresponding associated context of each historical observation may include, among other information, any information that is relevant to the specified task that the agent was performing in the past but that may not be obvious from the historical observation itself.

[0124] Typically, the system can repeatedly perform the following steps 404-410 of process 400 at each time step in a plurality of time steps to select a corresponding action (hereinafter referred to as the "current" action) to be performed by the agent at a corresponding state (hereinafter referred to as the "current" state) of the environment corresponding to the time step (hereinafter referred to as the "current" time step).

[0125] The system receives a current observation that characterizes the current state of the environment at the current time step (step 404). For example, the current observation may include an image, a video frame, a segment of audio data, a sentence of natural language, etc. In some of these examples, the observation may also include information derived from previous time steps, such as previous actions performed, rewards received at previous time steps, or both.

[0126] The system selects one or more historical observations from the plurality of historical observations (step 406). Figure 4 Explaining step 406 in more detail, Figure 4 Sub-steps 502 - 508 of step 406 are shown.

[0127] The system processes the current observation using an embedding neural network to generate an embedding, and then uses the embedding to generate a query vector (step 502). In some implementations, the embedding may be an activation output of an intermediate layer in the embedding neural network that has been generated based on processing the current observation, and the query vector is a projected embedding generated by applying a linear projection to the embedding.

[0128] The system selects one or more historical observations from the plurality of historical observations based on the query vector using a vector similarity search technique (step 504), such as based on vector similarity between the query vector and key vectors of the historical observations. The vector similarity may be an inner product, such as a dot product or a cosine similarity.

[0129] More specifically, this may include: for each historical observation stored in the retrieval data set, calculating the corresponding distance between the corresponding key vector of the historical observation and the query vector of the current observation, and then selecting one or more key vectors from the key vectors associated with multiple historical observations based on the corresponding distances of the one or more key vectors from the query vector of the current observation (step 506), for example, selecting the closest key vector.

[0130] This can be accomplished by using a maximum inner product search technique (MIPS) to compute the inner product of the query vector and the key vector to identify the query-key vector pair with the largest or maximum inner product (i.e., the closest distance, or in other words, the highest similarity). For example, the system identifies one or more key vectors that have a minimum distance, such as a Euclidean, Hamming, or other type of distance, to the query vector, or one or more key vectors that meet a distance threshold.

[0131] However, when a large number of historical observations, such as one million, ten million, one billion or more, are stored in the retrieval data set, the calculation of the inner product between the query vector and the corresponding key vector of the historical observation is slow and processor resource intensive. Therefore, some embodiments of the system use approximate nearest neighbor matching techniques, that is, instead of brute force methods, to achieve faster computation time while maintaining a high level of accuracy. For example, the system can use search space pruning, quantization of MIPS, or both. As a specific example of quantization technology, the system can use anisotropic quantization-based MIPS technology described in more detail at Ruiqi Guo, etal. "Accelerating large-scale inference with anisotropic vector quantization." International Conference on Machine Learning. PMLR, 2020. In anisotropic quantization-based MIPS, a quantized query is used to calculate the inner product, and the inner product has anisotropic weighting, which greatly penalizes errors parallel to the key compared to errors orthogonal to the key.

[0132] Then, the system selects corresponding historical observations from among a plurality of historical observations associated with the selected one or more key vectors in the retrieved data set as the selected one or more historical observations (step 508 ).

[0133] The system processes an encoder network input including (i) a current observation and (ii) a selected one or more historical observations and their respective associated contexts using an encoder neural network and according to parameter values ​​of the encoder neural network to generate a latent state representation of a current state of the environment (step 508).

[0134] The system uses the latent state representation to determine an action to be performed by the agent in response to the current observation (step 410). In some embodiments, the system processes the latent state representation using a policy neural network to generate a probability value for each action in a set of possible actions that can be performed by the agent, and then uses the probability values ​​to select an action to be performed by the agent, for example, by sampling actions according to their probability values, or by selecting the action with the highest probability value. In some other embodiments, the system performs a lookahead search that involves performing multiple planning iterations guided by the policy neural network, and then selecting an action based on the results of the lookahead search, i.e., after the last planning iteration.

[0135] Figure 6 Quantitative examples of the performance gains that can be achieved by an agent controlled using the reinforcement learning system described in this specification are shown. Figure 6 In the examples of , the environment is a simulated (or virtual) environment, and the agent is a computerized agent implemented as one or more computer programs that competes with a reference opponent agent to achieve a goal, where the actions in the action set are possible actions that can be performed by the computerized agent, and the goal is to win the competition with the reference opponent agent.

[0136] Figure 6 (a) shows the win rate against a reference opponent agent when the reinforcement learning system 100 selects actions to be performed by the agent based on different numbers (5, 10, 15, and 20) of retrieved historical observations, and the win rate against the reference opponent agent when a baseline non-retrieval agent control system (i.e., based on no retrieved historical observations) selects actions to be performed by the agent.

[0137] Figure 6 (b) shows the win rate as a function of the model size (defined as the number of times a model component, such as a convolution or attention block, is repeated throughout the policy neural network) of the retrieval-augmented policy neural network or the baseline off-policy neural network.

[0138] Understandably, the retrieval process improves the win rate against the same fixed reference opponent agent and also leads to improved performance compared to a baseline non-retrieval policy neural network of the same capacity.

[0139] This specification uses the term "configuration" in conjunction with system and computer program components. For a system of one or more computers to be configured to perform a particular operation or action, it is meant that the system has installed thereon software, firmware, hardware, or a combination thereof, which, in operation, causes the system to perform the operation or action. For one or more computer programs to be configured to perform a particular operation or action, it is meant that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0140] Embodiments of the subject matter and functional operations described in this specification may be implemented in digital electronic circuits, in tangibly embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in a combination of one or more of them. Embodiments of the subject matter described in this specification may be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium, for execution by a data processing device or for controlling the operation of a data processing device. A computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or additionally, program instructions may be encoded on an artificially generated propagation signal, such as a machine-generated electrical signal, optical signal, or electromagnetic signal, which is generated to encode information for transmission to a suitable receiver device for execution by a data processing device.

[0141] The term "data processing apparatus" refers to data processing hardware and covers all types of apparatus, devices and machines for processing data, including, for example, a programmable processor, a computer or multiple processors or computers. The apparatus may also be or further include special-purpose logic circuits, such as an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). In addition to the hardware, the apparatus may optionally include code that creates an execution environment for a computer program, for example, code constituting processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0142] A computer program, which may also be referred to or described as a program, software, software application, application, module, software module, script, or code, may be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it may be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program may be stored in a portion of a file that holds other programs or data, such as one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, such as files storing one or more modules, subroutines, or portions of code. A computer program may be deployed to execute on one computer located at one site or on multiple computers distributed across multiple sites and interconnected by a data communications network.

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

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

[0145] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuits (e.g., FPGAs or ASICs) or by a combination of special purpose logic circuits and one or more programmed computers.

[0146] Computers suitable for executing computer programs can be based on general or special microprocessors or both, or any other type of central processing unit. Typically, the central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The basic elements of a computer are a central processing unit for implementing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by or incorporated into a special logic circuit. Typically, the computer will also include or be operably coupled to receive data from or transmit data to or both of one or more large-capacity storage devices for storing data, such as magnetic disks, magneto-optical disks, or optical disks. However, the computer does not need to have such a device. In addition, the computer can be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, such as a universal serial bus (USB) flash drive, to name a few.

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

[0148] To provide interaction with a user, embodiments of the subject matter described in this specification may be implemented on a computer having a display device for displaying information to the user, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, 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 may also be used to provide interaction with the user; for example, the feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and input from the user may be received in any form, including acoustic, voice, or tactile input. In addition, the computer may interact with the user by sending documents to and receiving documents from a device used by the user; for example, by sending a web page to a web browser on a user's device in response to a request received from the web browser. In addition, the computer may interact with the user by sending a text message or other form of message to a personal device, such as a smart phone running a messaging application, and receiving a response message from the user in return.

[0149] The data processing apparatus for implementing the machine learning model may also include, for example, dedicated hardware accelerator units for processing common and computationally intensive parts of machine learning training or production—i.e., inference—workloads.

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

[0151] Embodiments of the subject matter described in this specification may be implemented in a computing system that includes a back-end component, such as a data server, or includes a middleware component, such as an application server, or includes a front-end component, such as a client computer with a graphical user interface, a web browser, or an application through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system may be interconnected by any form or medium of digital data communication, such as a communication network. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0152] The computing system may include a client and a server. The client and the server are usually remote from each other and usually interact through a communication network. The relationship between the client and the server is generated by computer programs running on the respective computers and having a client-server relationship with each other. In some embodiments, the server transmits data such as HTML pages to the user device, for example, for the purpose of displaying data to the user of the device interacting with the client and receiving user input from it. Data generated at the user device, such as the result of the user interaction, can be received from the device at the server.

[0153] Although this specification contains many specific implementation details, these should not be interpreted as limitations on the scope of any invention or on the scope that can be claimed, but rather as descriptions of features that can be specific to a particular embodiment of a particular invention. Certain features described in this specification in the context of separate embodiments may also be implemented in combination in a single embodiment. Conversely, the various features described in the context of a single embodiment may also be implemented in multiple embodiments separately or in any suitable sub-combination. In addition, although features may be described above as working in certain combinations and even initially claimed as such, in some cases, one or more features from the claimed combination may be removed from the combination, and the claimed combination may involve a sub-combination or a variation of the sub-combination.

[0154] Similarly, although operations are depicted in a particular order in the drawings and recited in the claims, this should not be understood as requiring that such operations be performed in the particular order shown or in a sequential order, or that all of the operations shown be performed to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous. In addition, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

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

Claims

1. A method for controlling a reinforcement learning agent in an environment to perform a task, the method comprising: maintaining a retrieval data set storing a plurality of historical observations and a corresponding associated context for each historical observation, wherein each historical observation represents a corresponding historical state of the environment and the corresponding associated context of each historical observation represents one or more actions performed in response to the historical observation; receiving a current observation characterizing a current state of the environment; Selecting one or more historical observations from the plurality of historical observations comprises: Generate a query vector for the current observation, including: processing the current observation using an embedding neural network to generate an embedding and processing the embedding to generate the query vector, and selecting the one or more historical observations from the plurality of historical observations using the query vector; and processing an encoder network input using an encoder neural network according to current values ​​of encoder network parameters to generate a latent state representation of the current state of the environment, the encoder network input comprising (i) the current observation and (ii) the selected one or more historical observations and their respective associated contexts; and An action to be performed by the agent in response to the current observation is determined using the latent state representation.

2. The method according to claim 1, wherein: The associated context of each historical observation includes one or both of the following: historical actions performed in response to the historical observation, or historical rewards received in response to the historical actions being performed.

3. The method according to any one of claims 1 to 2, wherein: The associated context of each historical observation further includes one or both of the following: a historical final observation characterizing a terminal state in a trajectory including the historical observation, a historical final reward received by the agent in the trajectory including the historical observation.

4. A method according to any one of claims 1 to 3, wherein Processing the embedding to generate the query vector includes: The embeddings are processed by using principal component analysis.

5. The method according to any one of claims 1 to 4, wherein: Each historical observation stored in the retrieval dataset is associated with a corresponding key vector generated by using the embedding neural network.

6. The method of claim 5, wherein selecting the one or more historical observations from the plurality of historical observations comprises: selecting one or more key vectors from among the key vectors associated with the plurality of historical observations based on: respective distances of the key vectors to the query vector of the current observation; and A corresponding historical observation is selected from the plurality of historical observations associated with the selected one or more key vectors as the selected one or more historical observations.

7. The method according to claim 6, wherein: Selecting the one or more key vectors comprises: One or more key vectors having a minimum Euclidean distance to the query vector of the current observation are selected.

8. The method according to any one of claims 1 to 7, wherein: Selecting the one or more historical observations from the plurality of historical observations using the query vector includes using an approximate nearest neighbor matching technique.

9. The method according to any one of claims 1 to 8, wherein: Selecting the one or more historical observations from the plurality of historical observations using the query vector includes using a vector similarity search technique.

10. The method according to claim 9, wherein: The vector similarity search technology includes a maximum inner product search technology.

11. The method according to any one of claims 9 to 10, wherein: The maximum inner product search technique includes a maximum inner product search technique based on anisotropic quantization.

12. The method according to any one of claims 1 to 11, wherein: Some or all of the plurality of historical observations stored in the retrieval dataset are generated as a result of the reinforcement learning agent performing one or more different tasks while interacting with the environment.

13. The method according to any one of claims 1 to 11, wherein: Some or all of the plurality of historical observations stored in the retrieval data set each have a different dimension, a different modality, or both than the current observation.

14. The method according to any one of claims 1 to 13, wherein: Using the latent state representation to determine the action to be performed by the agent in response to the current observation comprises: Perform a look-ahead search starting from the current state and searching a latent state representation space guided by a policy network output of a policy neural network, the policy neural network configured to process, at each time step in a plurality of time steps, a corresponding latent state representation for the time step to generate a corresponding policy network output for the time step.

15. The method according to claim 14, wherein: The look-ahead search includes a Monte Carlo Tree Search (MCTS) technique.

16. The method according to any one of claims 14 to 15, wherein: The policy network outputs include: a predicted policy output defining a distribution of scores over a set of possible actions; and A predicted value output defines a value representing a predicted long-term reward received by the agent as a result of the environment being in the current state.

17. The method according to any one of claims 1 to 13, wherein: Using the latent state representation to determine the action to be performed by the agent in response to the current observation comprises: The action to be performed by the agent is selected from the set of possible actions by using a policy network output of a policy neural network configured to process the latent state representation of the current state of the environment to generate the policy network output.

18. The method according to any one of claims 1 to 17, wherein: The task comprises one of: an object manipulation task or an environment navigation task.

19. The method according to any one of claims 1 to 18, wherein: The agent is a mechanical agent, the environment is a real-world environment, and the observations include data from one or more sensors configured to sense the real-world environment.

20. The method according to any one of claims 1 to 19, wherein: The encoder neural network comprises: a first subnetwork configured to process the current observation to generate a first encoded representation; a second sub-network configured to process each of the selected one or more historical observations and their respective associated contexts to generate a respective second encoded representation; and A third sub-network is configured to process the first encoded representation and the corresponding second encoded representation to generate the latent state representation.

21. The method according to any one of claims 1 to 20, wherein: The encoder neural network has been trained by reinforcement learning on a training dataset storing a plurality of trajectories generated as a result of the reinforcement learning agent interacting with the environment.

22. The method according to claim 21, wherein: The retrieval data set is the same as the training data set.

23. The method according to claim 21, wherein: The retrieval data set is different from the training data set and is larger than the training data set.

24. The method according to any one of claims 21 to 23, wherein: The encoder neural network has been trained by model-based reinforcement learning.

25. One or more computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the respective methods of any one of claims 1 to 24.

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

Images