Generating environment models using in-context adaptation and exploration

A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations for controlling an agent interacting with an environment to perform a task, and wherein the operations comprise, at a current time step of multiple time steps:

• maintaining context data that includes, for each of one or more previous time steps that precede the current time step, (i) data identifying a previous action performed by the agent at the previous time step and (ii) a previous observation characterizing a previous state that the environment transitioned into as a result of the agent performing the previous action;
• receiving a current observation characterizing a current state of the environment;
• generating a current graph model that represents the environment, wherein the current graph model comprises vertices that represent states of the environment and edges connecting the vertices, wherein an edge between a first vertex and a second vertex in the graph model indicates it is possible that the environment will transition from a state represented by the first vertex into a state represented by the second vertex as a result of one of a possible set of actions performed by the agent, and wherein generating the current graph model comprises:
• processing a first Transformer network input that includes the context data using a Transformer neural network to generate a probability distribution over a possible set of edges that can be included in the current graph model;
• and
• selecting, in accordance with the probability distribution over the possible set of edges, a subset of the edges to be included in the current graph model;
• selecting, from the possible set of actions and using the current graph model, a current action to be performed by the agent in response to the current observation;
• controlling the agent to perform the selected current action to cause the environment to transition from the current state into a new state;
• and
• updating the context data to include (i) data identifying the selected current action and (ii) a new observation characterizing the new state of the environment.

A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform operations for controlling an agent interacting with an environment to perform a task, and wherein the operations comprise, at a current time step of multiple time steps:

• maintaining context data that includes, for each of one or more previous time steps that precede the current time step, (i) data identifying a previous action performed by the agent at the previous time step and (ii) a previous observation characterizing a previous state that the environment transitioned into as a result of the agent performing the previous action;
• receiving a current observation characterizing a current state of the environment;
• generating a current graph model that represents the environment, wherein the current graph model comprises vertices that represent states of the environment and edges connecting the vertices, wherein an edge between a first vertex and a second vertex in the graph model indicates it is possible that the environment will transition from a state represented by the first vertex into a state represented by the second vertex as a result of one of a possible set of actions performed by the agent, and wherein generating the current graph model comprises:
• processing a first Transformer network input that includes the context data using a Transformer neural network to generate a probability distribution over a possible set of edges that can be included in the current graph model;
• and
• selecting, in accordance with the probability distribution over the possible set of edges, a subset of the edges to be included in the current graph model;
• selecting, from the possible set of actions and using the current graph model, a current action to be performed by the agent in response to the current observation;
• controlling the agent to perform the selected current action to cause the environment to transition from the current state into a new state;
• and
• updating the context data to include (i) data identifying the selected current action and (ii) a new observation characterizing the new state of the environment.