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 training a policy neural network having a plurality of policy neural network parameters used to select actions to be performed by an agent to control the agent in an environment, the operations comprising:
- obtaining an initial observation that characterizes an initial state of the environment;
- processing at least the initial observation by using a representation neural network having a plurality of representation neural network parameters to generate an initial latent representation that characterizes the initial state of the environment, wherein the initial latent representation is in a latent space;
- generating a trajectory of latent representations that comprises a sequence of multiple latent representations that are in the latent space and that each characterize a different state of the environment, wherein the sequence of multiple latent representations start with the initial latent representation that characterizes the initial state of the environment, and wherein generating the trajectory of latent representations comprises, at each subsequent step that is after a beginning step in a plurality of steps:
- processing a current latent representation that is generated by using a transition neural network and without processing any observation that characterizes any state of the environment using the policy neural network and in accordance with current values of the plurality of policy neural network parameters to generate an action selection output;
- selecting, using the action selection output, an action that would be performed by the agent when the environment is in a current state characterized by the current latent representation;
- and
- processing the current latent representation that is generated by using the transition neural network and the action that is selected using the policy neural network and without processing any observation that characterizes any state of the environment using the transition neural network having a plurality of transition neural network parameters to generate an updated latent representation that is in the latent space and that characterizes a next state that the environment would transition into if the agent performed the selected action when the environment is in the current state characterized by the current latent representation;
- for each latent representation in the sequence of multiple latent representations included in the trajectory that are in the latent space:
- determining a predicted reward;
- and
- processing the latent representation and without processing any observation that characterizes any state of the environment using a value neural network having a plurality of value neural network parameters and in accordance with current values of the plurality of value neural network parameters to generate a predicted state value that represents a time-adjusted sum of predicted rewards that would be received by the agent by selecting actions using action selection outputs starting from a state characterized by the latent representation;
- determining a target state value for each latent representation in the trajectory from the predicted rewards and the predicted state values for the latent representations in the trajectory;
- training, based on the predicted state values that have been generated by using the value neural network from the latent representations in the latent space, the policy neural network to update the current values of the plurality of policy neural network parameters, wherein training the policy neural network comprises determining an update to the current values of the plurality of policy neural network parameters that encourages the policy neural network to generate action selection outputs that result in actions being selected that improve the predicted state values;
- and
- training, based on the predicted state values that have been generated by using the value neural network from the latent representations in the latent space, the value neural network to update the current values of the plurality of value neural network parameters, wherein training the value neural network comprises determining an update to the current values of the plurality of value neural network parameters by optimizing a value objective function that measures a difference between the predicted state values and the target state values.