Training an unsupervised memory-based prediction system to learn compressed representations of an environment

A real-world system comprising one or more processors and one or more storage devices storing instructions that when executed by the one or more processors cause the one or more processors to perform operations for training a memory-based prediction neural network system using a predictive modeling process, wherein the memory-based prediction neural network system includes an encoder and a representation system and has a plurality of memory-based prediction parameters and is configured to:

• receive, by the memory-based prediction neural network system, an input observation characterizing a state of an environment interacted with by an agent;
• and
• process, by the memory-based prediction neural network system, the input observation and data read from a memory to update data stored in the memory and to generate a latent representation of the state of the environment;
• the operations comprising:
• performing, by the memory-based prediction neural network system, for each of a plurality of time steps:
• receiving a current observation for the time step characterizing a current state of the environment being interacted with by the agent at the time step;
• obtaining the data read from the memory, comprising:
• processing one or more latent representations generated at a previous time step in accordance with current values of the memory-based prediction parameters to generate one or more read keys;
• and
• generating one or more readout vectors based on a first measure of similarity between:
• (i) the one or more read keys derived from the one or more latent representations generated at the previous time step, and (ii) data stored in the memory,
• wherein the one or more readout vectors define the data read from the memory;
• processing the current observation for the time step and data read from the memory in accordance with the current values of the memory-based prediction parameters to:
• (i) update the data stored in the memory, and (ii) generate a latent representation of the current state of the environment as of the time step, comprising:
• processing the current observation using the encoder and in accordance with the current values of the memory-based prediction parameters to generate an encoded representation of the current observation;
• processing:
• (i) the encoded representation of the current observation, and (ii) the data read from the memory, using the representation system and in accordance with the current values of the memory-based prediction parameters, to generate the latent representation of the current state of the environment as of the time step;
• and
• writing the latent representation of the current state of the environment as of the time step to the memory;
• and
• generating, using the latent representation of the current state of the environment as of the time step and in accordance with the current values of the memory-based prediction parameters, data characterizing a predicted return, wherein the predicted return is a prediction for a cumulative measure of rewards that will be received by the agent as a result of interactions with the environment to perform a task after the current observation for the time step is received;
• determining, for one or more of the time steps, an actual return for the time step that is a cumulative measure of rewards received by the agent as a result of interactions with the environment to perform the task after the current observation for the time step is received;
• determining a gradient based at least in part on, for one or more of the time steps, an error between:
• (i) the data characterizing the predicted return for the time step, and (ii) the actual return for the time step;
• and
• adjusting the current values of the memory-based prediction parameters using the gradient.

One or more non-transitory computer storage media storing instructions that when executed by one or more processors cause the one or more processors to perform operations for training a memory-based prediction neural network system using a predictive modeling process, wherein the memory-based prediction neural network system includes an encoder and a representation system and has a plurality of memory-based prediction parameters and is configured to:

• receive, by the memory-based prediction neural network system, an input observation characterizing a state of an environment interacted with by an agent;
• and
• process, by the memory-based prediction neural network system, the input observation and data read from a memory to update data stored in the memory and to generate a latent representation of the state of the environment;
• the operations comprising:
• performing, by the memory-based prediction neural network system, for each of a plurality of time steps:
• receiving a current observation for the time step characterizing a current state of the environment being interacted with by the agent at the time step;
• obtaining the data read from the memory, comprising:
• processing one or more latent representations generated at a previous time step in accordance with current values of the memory-based prediction parameters to generate one or more read keys;
• and
• generating one or more readout vectors based on a first measure of similarity between:
• (i) the one or more read keys derived from the one or more latent representations generated at the previous time step, and (ii) data stored in the memory,
• wherein the one or more readout vectors define the data read from the memory;
• processing the current observation for the time step and data read from the memory in accordance with the current values of the memory-based prediction parameters to:
• (i) update the data stored in the memory, and (ii) generate a latent representation of the current state of the environment as of the time step, comprising:
• processing the current observation using the encoder and in accordance with the current values of the memory-based prediction parameters to generate an encoded representation of the current observation;
• processing:
• (i) the encoded representation of the current observation, and (ii) the data read from the memory, using the representation system and in accordance with the current values of the memory-based prediction parameters, to generate the latent representation of the current state of the environment as of the time step;
• and
• writing the latent representation of the current state of the environment as of the time step to the memory;
• and
• generating, using the latent representation of the current state of the environment as of the time step and in accordance with the current values of the memory-based prediction parameters, data characterizing a predicted return, wherein the predicted return is a prediction for a cumulative measure of rewards that will be received by the agent as a result of interactions with the environment to perform a task after the current observation for the time step is received;
• determining, for one or more of the time steps, an actual return for the time step that is a cumulative measure of rewards received by the agent as a result of interactions with the environment to perform the task after the current observation for the time step is received;
• determining a gradient based at least in part on, for one or more of the time steps, an error between:
• (i) the data characterizing the predicted return for the time step, and (ii) the actual return for the time step;
• and
• adjusting the current values of the memory-based prediction parameters using the gradient.