Action selection neural network training using imitation learning in latent space

1. A method for training an action selection policy neural network, wherein the action selection policy neural network has a plurality of action selection policy neural network parameters, wherein the action selection policy neural network is configured to process an observation characterizing a state of an environment in accordance with values of the action selection policy neural network parameters to generate an action selection policy output, wherein the action selection policy output is used to select an action to be performed by an agent interacting with an environment, the method comprising:

obtaining an observation characterizing a state of the environment subsequent to the agent performing a selected action;

generating a latent representation of the observation, comprising:

(i) processing the observation using an encoder neural network to generate the latent representation of the observation, wherein the encoder neural network has been trained to process a given observation to generate a latent representation of the given observation that is predictive of latent representations generated by the encoder neural network of one or more other observations that are after the given observation in a sequence of observations; or

(ii) processing the observation using the action selection policy neural network to generate the latent representation of the observation as an intermediate output of the action selection policy neural network;

processing the latent representation of the observation using a discriminator neural network to generate an imitation score,

wherein the imitation score characterizes a likelihood that the observation is included in an expert trajectory,

wherein an expert trajectory is a sequence of observations characterizing respective states of the environment while a given agent interacts with the environment by performing a sequence of actions that accomplish a particular task;

determining a reward from the imitation score; and

adjusting the current values of the action selection policy neural network parameters based on the reward using a reinforcement learning training technique.

2. The method of claim 1 , further comprising:

obtaining a task reward characterizing a progress of the agent towards accomplishing the particular task; and

determining the reward from the task reward in addition to the imitation score.

3. The method of claim 2 , wherein determining the reward from the task reward in addition to the imitation score comprises:

determining the reward as a weighted linear combination of the task reward and the imitation score.

4. The method of claim 1 , further comprising:

obtaining an expert observation from an expert trajectory in a set of multiple expert trajectories;

processing the expert observation using the encoder neural network to generate a latent representation of the expert observation;

processing the latent representation of the expert observation using the discriminator neural network to generate a given imitation score;

determining a loss based on a difference between the given imitation score and a target imitation score which indicates that the expert observation is from an expert trajectory; and

adjusting current values of discriminator neural network parameters based on the loss.

5. The method of claim 1 , further comprising:

determining a loss based on a difference between the imitation score and a target imitation score which indicates that the observation is not from an expert trajectory; and

adjusting current values of discriminator neural network parameters based on the loss.

6. The method of claim 1 , wherein the observation comprises an image.

7. The method of claim 1 , wherein the encoder neural network is trained on training data comprising a plurality of expert trajectories.

8. The method of claim 1 , further comprising:

processing the latent representation of the observation using a second discriminator neural network to generate a second imitation score,

wherein the second imitation score characterizes a likelihood that the observation is: (i) included in an expert trajectory, and (ii) within a threshold number of observations of a last observation of the expert trajectory; and

determining the reward from the second imitation score in addition to the imitation score.

9. The method of claim 8 , further comprising:

obtaining an expert observation from an expert trajectory, wherein the expert observation is within the threshold number of observations of the last observation of the expert trajectory;

processing the expert observation using the encoder neural network to generate a latent representation of the expert observation;

processing the latent representation of the expert observation using the second discriminator neural network to generate a given second imitation score;

determining a loss based on a difference between the given second imitation score and a target imitation score which indicates that the expert observation is from an expert trajectory and is within the threshold number of observations of the last observation in the expert trajectory; and

adjusting current values of second discriminator neural network parameters based on the loss.

10. The method of claim 8 , further comprising:

processing the latent representation of the observation using the second discriminator neural network to generate a particular second imitation score;

determining a loss based on a difference between the particular second imitation score and a target imitation score which indicates that the observation is not from within the threshold number of observations of the last observation in an expert trajectory; and

adjusting current values of second discriminator neural network parameters based on the loss.

11. The method of claim 1 , wherein the discriminator neural network has a single neural network layer.

12. The method of claim 1 , wherein obtaining an observation characterizing a state of the environment subsequent to the agent performing a selected action comprises:

sampling the observation from a replay buffer.

13. The method of claim 1 , wherein the reinforcement learning training technique is a deep deterministic policy gradient training technique.

14. The method of claim 1 , wherein the action selection policy neural network comprises a Q network and a policy network.

15. 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 training an action selection policy neural network, wherein the action selection policy neural network has a plurality of action selection policy neural network parameters, wherein the action selection policy neural network is configured to process an observation characterizing a state of an environment in accordance with values of the action selection policy neural network parameters to generate an action selection policy output, wherein the action selection policy output is used to select an action to be performed by an agent interacting with an environment, the operations comprising:

obtaining an observation characterizing a state of the environment subsequent to the agent performing a selected action;

generating a latent representation of the observation, comprising:

(i) processing the observation using an encoder neural network to generate the latent representation of the observation, wherein the encoder neural network has been trained to process a given observation to generate a latent representation of the given observation that is predictive of latent representations generated by the encoder neural network of one or more other observations that are after the given observation in a sequence of observations; or

(ii) processing the observation using the action selection policy neural network to generate the latent representation of the observation as an intermediate output of the action selection policy neural network;

processing the latent representation of the observation using a discriminator neural network to generate an imitation score,

wherein the imitation score characterizes a likelihood that the observation is included in an expert trajectory,

wherein an expert trajectory is a sequence of observations characterizing respective states of the environment while a given agent interacts with the environment by performing a sequence of actions that accomplish a particular task;

determining a reward from the imitation score;

adjusting the current values of the action selection policy neural network parameters based on the reward using a reinforcement learning training technique.

16. The system of claim 15 , wherein the operations further comprise:

obtaining a task reward characterizing a progress of the agent towards accomplishing the particular task; and

determining the reward from the task reward in addition to the imitation score.

17. The system of claim 16 , wherein determining the reward from the task reward in addition to the imitation score comprises:

determining the reward as a weighted linear combination of the task reward and the imitation score.

18. 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 training an action selection policy neural network, wherein the action selection policy neural network has a plurality of action selection policy neural network parameters, wherein the action selection policy neural network is configured to process an observation characterizing a state of an environment in accordance with values of the action selection policy neural network parameters to generate an action selection policy output, wherein the action selection policy output is used to select an action to be performed by an agent interacting with an environment, the operations comprising:

obtaining an observation characterizing a state of the environment subsequent to the agent performing a selected action;

generating a latent representation of the observation, comprising:

(i) processing the observation using an encoder neural network to generate the latent representation of the observation, wherein the encoder neural network has been trained to process a given observation to generate a latent representation of the given observation that is predictive of latent representations generated by the encoder neural network of one or more other observations that are after the given observation in a sequence of observations; or

(ii) processing the observation using the action selection policy neural network to generate the latent representation of the observation as an intermediate output of the action selection policy neural network;

processing the latent representation of the observation using a discriminator neural network to generate an imitation score,

wherein the imitation score characterizes a likelihood that the observation is included in an expert trajectory,

wherein an expert trajectory is a sequence of observations characterizing respective states of the environment while a given agent interacts with the environment by performing a sequence of actions that accomplish a particular task;

determining a reward from the imitation score;

adjusting the current values of the action selection policy neural network parameters based on the reward using a reinforcement learning training technique.

19. The non-transitory computer storage media of claim 18 , wherein the operations further comprise:

obtaining a task reward characterizing a progress of the agent towards accomplishing the particular task; and

determining the reward from the task reward in addition to the imitation score.

20. The non-transitory computer storage media of claim 19 , wherein determining the reward from the task reward in addition to the imitation score comprises:

determining the reward as a weighted linear combination of the task reward and the imitation score.