Generating spatial embeddings by integrating agent motion and optimizing a predictive objective
Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training a spatial embedding neural network that is configured to process data characterizing motion of an agent that is interacting with an environment to generate spatial embeddings. In one aspect, a method comprises: processing data characterizing the motion of the agent in the environment at the current time step using a spatial embedding neural network to generate a current spatial embedding for the current time step; determining a predicted score and a target score for each of a plurality of slots in an external memory, wherein each slot stores: (i) a representation of an observation characterizing a state of the environment, and (ii) a spatial embedding; and determining an update to values of the set of spatial embedding neural network parameters based on an error between the predicted scores and the target scores.
1 . A method performed by one or more computers for selecting actions to be performed by an agent to efficiently navigate an environment to perform a task, the method comprising:
controlling the agent to interact with the environment based on one or more spatial embeddings that encode spatial information for the environment, including repeatedly performing operations comprising:
processing current agent motion data using a spatial embedding neural network to generate a current spatial embedding; and
processing at least the current spatial embedding using an action selection neural network to generate a current action selection output; and
selecting a current action to be performed by the agent based on the current action selection output;
wherein the spatial embedding neural network has been trained using a training technique comprising, for each of a plurality of time steps:
processing agent motion data characterizing one or more of: a speed of the agent at the time step, an angular velocity of the agent at the time step, or a translational velocity of the agent at the time step, using the spatial embedding neural network to generate a current spatial embedding for the time step;
determining a predicted score and a target score for each of a plurality of slots in an external memory, wherein each slot corresponds to a respective previous time step and stores: (i) a representation of an observation comprising an image characterizing a state of the environment for the respective previous time step, and (ii) a spatial embedding for the respective previous time step generated based on agent motion data for the respective previous time step,
wherein the predicted score for each slot measures a similarity between: (i) the current spatial embedding, and (ii) the spatial embedding corresponding to the slot,
wherein the target score for each slot measures a similarity between: (i) a current observation comprising a current image characterizing the state of the environment at the time step, and (ii) the observation comprising the image corresponding to the slot; and
determining an update to values of a set of spatial embedding neural network parameters based on an error between the predicted scores and the target scores.
2 . The method of claim 1 , wherein the image is captured from a perspective of the agent at the time step.
3 . The method of claim 1 , wherein determining the target score for each slot in the external memory comprises:
obtaining respective embeddings of the current observation characterizing the current state of the environment and the observation corresponding to the slot; and
determining the target score based on a similarity measure between: (i) the embedding of the current observation characterizing the current state of the environment, and (ii) the embedding of the observation corresponding to the slot.
4 . The method of claim 3 , wherein obtaining the embedding of the current observation comprises processing the current observation using an embedding neural network.
5 . The method of claim 1 , wherein the error between the predicted scores and the target scores comprises a cross-entropy error between the predicted scores and the target scores.
6 . The method of claim 1 , further comprising determining an update to the spatial embeddings stored in the external memory based on the error between the predicted scores and the target scores.
7 . The method of claim 1 , wherein the spatial embedding neural network does not process the current observation to generate the current spatial embedding for the time step.
8 . The method of claim 1 , further comprising storing a representation of the current observation and the current spatial embedding in a slot in the external memory.
9 . The method of claim 1 , wherein the action selection output comprises a respective score for each action in a predetermined set of actions.
10 . The method of claim 9 , wherein selecting the current action to be performed by the agent based on the current action selection output comprises selecting an action having a highest score.
11 . The method of claim 1 , wherein the action selection neural network is configured to additionally processes a goal spatial embedding that was generated by the spatial embedding neural network at a time step when the agent was located in a goal location in the environment.
12 . The method of claim 1 , wherein the spatial embedding neural network comprises a recurrent neural network, and wherein generating the current spatial embedding for the time step comprises:
processing: (i) the agent motion data characterizing one or more of: a speed of the agent at the time step, an angular velocity of the agent at the time step, or a translational velocity of the agent at the time step, and (ii) an output of the spatial embedding neural network at a preceding time step, using the spatial embedding neural network to update a hidden state of the spatial embedding neural network, wherein the updated hidden state defines the current spatial embedding.
13 . The method of claim 12 , further comprising determining an output of the spatial embedding neural network for the time step, comprising:
identifying the updated hidden state of the spatial embedding neural network as the output of the spatial embedding neural network for the time step.
14 . The method of claim 12 , further comprising determining an output of the spatial embedding neural network for the time step, comprising:
determining a respective weight value for each slot in the external memory that characterizes a similarity between: (i) the current observation characterizing the state of the environment at the time step, and (ii) the observation corresponding to the slot;
determining a correction embedding as a linear combination of the spatial embeddings corresponding to the slots in the external memory, wherein each spatial embedding is weighted by the respective weight value; and
determining the output of the spatial embedding neural network based on: (i) the updated hidden state of the spatial embedding neural network, and (ii) the correction embedding.
15 . The method of claim 1 , further comprising:
processing agent motion data characterizing one or more of: a speed of the agent at the time step, an angular velocity of the agent at the time step, or a translational velocity of the agent at the time step using an integrated embedding neural network having a set of integrated embedding neural network parameters to generate a current integrated embedding for the time step;
determining a predicted score and a target score for each of a plurality of slots in an additional external memory, wherein each slot stores: (i) a spatial embedding, and (ii) an integrated embedding,
wherein the predicted score for each slot measures a similarity between: (i) the current integrated embedding, and (ii) the integrated embedding corresponding to the slot,
wherein the target score for each slot measures a similarity between: (i) the current spatial embedding, and (ii) the spatial embedding corresponding to the slot; and
determining an update to values of the set of integrated embedding neural network parameters based on an error between the predicted scores and the target scores for the slots in the additional external memory.
16 . 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 to efficiently navigate an environment to perform a task, the operations comprising:
controlling the agent to interact with the environment based on one or more spatial embeddings that encode spatial information for the environment, including repeatedly performing operations comprising:
processing current agent motion data using a spatial embedding neural network to generate a current spatial embedding; and
processing at least the current spatial embedding using an action selection neural network to generate a current action selection output; and
selecting a current action to be performed by the agent based on the current action selection output;
wherein the spatial embedding neural network has been trained using a training technique comprising, for each of a plurality of time steps:
processing agent motion data characterizing one or more of: a speed of the agent at the time step, an angular velocity of the agent at the time step, or a translational velocity of the agent at the time step, using the spatial embedding neural network to generate a current spatial embedding for the time step;
determining a predicted score and a target score for each of a plurality of slots in an external memory, wherein each slot corresponds to a respective previous time step and stores: (i) a representation of an observation comprising an image characterizing a state of the environment for the respective previous time step, and (ii) a spatial embedding for the respective previous time step generated based on agent motion data for the respective previous time step,
wherein the predicted score for each slot measures a similarity between: (i) the current spatial embedding, and (ii) the spatial embedding corresponding to the slot,
wherein the target score for each slot measures a similarity between: (i) a current observation comprising a current image characterizing the state of the environment at the time step, and (ii) the observation comprising the image corresponding to the slot; and
determining an update to values of a set of spatial embedding neural network parameters based on an error between the predicted scores and the target scores.
17 . 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 to efficiently navigate an environment to perform a task, the operations comprising:
controlling the agent to interact with the environment based on one or more spatial embeddings that encode spatial information for the environment, including repeatedly performing operations comprising:
processing current agent motion data using a spatial embedding neural network to generate a current spatial embedding; and
processing at least the current spatial embedding using an action selection neural network to generate a current action selection output; and
selecting a current action to be performed by the agent based on the current action selection output;
wherein the spatial embedding neural network has been trained using a training technique comprising, for each of a plurality of time steps:
processing agent motion data characterizing one or more of: a speed of the agent at the time step, an angular velocity of the agent at the time step, or a translational velocity of the agent at the time step, using the spatial embedding neural network to generate a current spatial embedding for the time step;
determining a predicted score and a target score for each of a plurality of slots in an external memory, wherein each slot corresponds to a respective previous time step and stores: (i) a representation of an observation comprising an image characterizing a state of the environment for the respective previous time step, and (ii) a spatial embedding for the respective previous time step generated based on agent motion data for the respective previous time step,
wherein the predicted score for each slot measures a similarity between: (i) the current spatial embedding, and (ii) the spatial embedding corresponding to the slot,
wherein the target score for each slot measures a similarity between: (i) a current observation comprising a current image characterizing the state of the environment at the time step, and (ii) the observation comprising the image corresponding to the slot; and
determining an update to values of a set of spatial embedding neural network parameters based on an error between the predicted scores and the target scores.
18 . The system of claim 17 , wherein the image is captured from a perspective of the agent at the time step.
19 . The system of claim 17 , wherein determining the target score for each slot in the external memory comprises:
obtaining respective embeddings of the current observation characterizing the current state of the environment and the observation corresponding to the slot; and
determining the target score based on a similarity measure between: (i) the embedding of the current observation characterizing the current state of the environment, and (ii) the embedding of the observation corresponding to the slot.