Generating implicit plans for accomplishing goals in an environment using attention operations over planning embeddings
Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for selecting actions to be performed by an agent interacting with an environment to accomplish a goal. In one aspect, a method comprises: generating a respective planning embedding corresponding to each of multiple experience tuples in an external memory, wherein each experience tuple characterizes interaction of the agent with the environment at a respective previous time step; processing the planning embeddings using a planning neural network to generate an implicit plan for accomplishing the goal; and selecting the action to be performed by the agent at the time step using the implicit plan.
1 . A method performed by one or more data processing apparatus for selecting actions to be performed by an agent interacting with an environment to accomplish a goal, the method comprising:
receiving data identifying a task to be performed by the agent, wherein the task comprises causing a state of the environment to transition into a goal state;
generating a representation of the goal state;
generating a respective planning embedding corresponding to each of multiple experience tuples in an external memory, wherein each experience tuple characterizes an interaction of the agent with the environment at a respective previous time step, wherein generating the planning embeddings for the experience tuples comprises:
appending the same generated representation of the goal state of the environment to each of the experience tuples;
generating an implicit plan for accomplishing the goal by processing, using a planning neural network, (i) the planning embeddings, and (ii) a current observation characterizing a current state of the environment; and
selecting the action to be performed by the agent at the time step using the implicit plan, wherein:
(i) the agent is a mechanical agent, and the actions include control signals for actions taken by the mechanical agent in the environment to accomplish the goal, or
(ii) the environment includes items of equipment in a plant or facility and the agent controls actions in the environment, the actions controlling or imposing operating conditions on the items of equipment of the plant/facility, or
(iii) the agent manages distribution of tasks across computing resources in the environment, and the actions include assigning tasks to particular computing resources.
2 . The method of claim 1 , wherein processing the planning embeddings using a planning neural network to generate an implicit plan for accomplishing the goal comprises:
iteratively updating the planning embeddings, comprising, at each of a plurality of iterations, updating each planning embedding using attention over the planning embeddings; and
generating the implicit plan using the planning embeddings.
3 . The method of claim 2 , further comprising generating a respective static embedding corresponding to each of multiple experience tuples in the external memory;
wherein iteratively updating the planning embeddings further comprises, at each of the plurality of iterations:
updating each planning embedding using attention over the static embeddings.
4 . The method of claim 3 , wherein:
generating a respective planning embedding corresponding to each of multiple experience tuples in the external memory comprises:
generating a respective planning embedding only for experience tuples in the external memory characterizing interaction of the agent with the environment over a predetermined number of most recent time steps; and
generating a respective static embedding corresponding to each of multiple experience tuples in the external memory comprises:
generating a respective static embedding for every experience tuple in the external memory.
5 . The method of claim 2 , wherein updating each planning embedding using attention over the planning embeddings comprises:
processing the planning embeddings using a residual neural network block that is configured to apply a sequence of operations to the planning embeddings, the sequence of operations comprising: (i) a layer normalization operation, (ii) an attention operation, and (iii) a linear projection operation.
6 . The method of claim 5 , wherein the attention operation comprises a multi-head key-query-value attention operation over the planning embeddings.
7 . The method of claim 2 , wherein generating the implicit plan using the planning embeddings comprises:
generating the implicit plan based on: (i) the planning embeddings, and (ii) a current observation characterizing a current state of the environment.
8 . The method of claim 7 , wherein generating the implicit plan based on: (i) the planning embeddings, and (ii) a current observation characterizing a current state of the environment comprises:
for each planning embedding:
appending a representation of the current observation to the planning embedding to generate a combined embedding; and
processing the combined embedding by one or more neural network layers to generate a transformed embedding; and
generating the implicit plan based on the transformed embeddings.
9 . The method of claim 8 , wherein generating the implicit plan based on the transformed embeddings comprises:
generating the implicit plan by applying a pooling operation to the transformed embeddings.
10 . The method of claim 9 , wherein the pooling operation is a feature-wise max pooling operation.
11 . The method of claim 1 , wherein selecting the action to be performed by the agent at the time step using the implicit plan comprises:
processing an input comprising the implicit plan using an action selection neural network to generate an action selection output; and
selecting the action based on the action selection output.
12 . The method of claim 11 , wherein the action selection output comprises a respective score for each action in a set of possible actions that can be performed by the agent, and selecting the action based on the action selection output comprises sampling an action in accordance with the action scores.
13 . The method of claim 11 , wherein the action selection neural network and the planning neural network are trained using reinforcement learning techniques to maximize a cumulative measure of reward received by the agent by interacting with the environment.
14 . The method of claim 1 , wherein each experience tuple comprises: (i) a representation of an observation characterizing a state of the environment at the respective previous time step, (ii) a representation of an action performed by the agent at the respective previous time step, and (iii) a representation of an observation characterizing the state of the environment after the agent performed the action at the respective previous time step.
15 . The method of claim 1 , further comprising, after selecting the action to be performed by the agent at the time step using the implicit plan, storing an experience tuple in the external memory that characterizes interaction of the agent with the environment at a current time step.
16 . A system comprising:
one or more computers; and
one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations for selecting actions to be performed by an agent interacting with an environment to accomplish a goal, the operations comprising:
receiving data identifying a task to be performed by the agent, wherein the task comprises causing a state of the environment to transition into a goal state;
generating a representation of the goal state;
generating a respective planning embedding corresponding to each of multiple experience tuples in an external memory, wherein each experience tuple characterizes an interaction of the agent with the environment at a respective previous time step, wherein generating the planning embeddings for the experience tuples comprises:
appending the same generated representation of the goal state of the environment to each of the experience tuples;
generating an implicit plan for accomplishing the goal by processing, using a planning neural network, (i) the planning embeddings, and (ii) a current observation characterizing a current state of the environment; and
selecting the action to be performed by the agent at the time step using the implicit plan, wherein:
(i) the agent is a mechanical agent, and the actions include control signals for actions taken by the mechanical agent in the environment to accomplish the goal, or
(ii) the environment includes items of equipment in a plant or facility and the agent controls actions in the environment, the actions controlling or imposing operating conditions on the items of equipment of the plant/facility, or
(iii) the agent manages distribution of tasks across computing resources in the environment, and the actions include assigning tasks to particular computing resources.
17 . One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for selecting actions to be performed by an agent interacting with an environment to accomplish a goal, the operations comprising:
receiving data identifying a task to be performed by the agent, wherein the task comprises causing a state of the environment to transition into a goal state;
generating a representation of the goal state;
generating a respective planning embedding corresponding to each of multiple experience tuples in an external memory, wherein each experience tuple characterizes an interaction of the agent with the environment at a respective previous time step, wherein generating the planning embeddings for the experience tuples comprises:
appending the same generated representation of the goal state of the environment to each of the experience tuples;
generating an implicit plan for accomplishing the goal by processing, using a planning neural network, (i) the planning embeddings, and (ii) a current observation characterizing a current state of the environment; and
selecting the action to be performed by the agent at the time step using the implicit plan, wherein:
(i) the agent is a mechanical agent, and the actions include control signals for actions taken by the mechanical agent in the environment to accomplish the goal, or
(ii) the environment includes items of equipment in a plant or facility and the agent controls actions in the environment, the actions controlling or imposing operating conditions on the items of equipment of the plant/facility, or
(iii) the agent manages distribution of tasks across computing resources in the environment, and the actions include assigning tasks to particular computing resources.
18 . The non-transitory computer storage media of claim 17 , wherein processing the planning embeddings using a planning neural network to generate an implicit plan for accomplishing the goal comprises:
iteratively updating the planning embeddings, comprising, at each of a plurality of iterations, updating each planning embedding using attention over the planning embeddings; and
generating the implicit plan using the planning embeddings.
19 . The non-transitory computer storage media of claim 17 , the operations further comprising generating a respective static embedding corresponding to each of multiple experience tuples in the external memory;
wherein iteratively updating the planning embeddings further comprises, at each of the plurality of iterations:
updating each planning embedding using attention over the static embeddings.
20 . The non-transitory computer storage media of claim 19 , the operations further comprising generating a respective static embedding corresponding to each of multiple experience tuples in the external memory;
wherein iteratively updating the planning embeddings further comprises, at each of the plurality of iterations:
updating each planning embedding using attention over the static embeddings.