Learning policies using sparse and underspecified rewards

1 . A method of learning a policy having a plurality of policy parameters configured to receive a context input for an agent to perform a task and to generate, from the context input and in accordance with the policy parameters, a trajectory of one or more actions for performing the task, wherein the one or more actions comprise control inputs for controlling the agent, the method comprising:

obtaining a plurality of training context inputs and a plurality of validation context inputs;

generating successful trajectory data that identifies a plurality of successful trajectories, each successful trajectory corresponding to a respective context input from either the training context inputs or the validation context inputs and being a trajectory for which a received task reward value indicates that the task was successfully completed by performing the trajectory in response to the corresponding context input; and

training the policy jointly with an auxiliary reward function having a plurality of auxiliary reward parameters, the auxiliary reward function being configured to map, in accordance with the auxiliary reward parameters, trajectory features of at least a trajectory to an auxiliary reward value that indicates how well the trajectory performed the task in response to the context input, and the training comprising repeatedly performing at least the following operations:

obtaining, from the successful trajectory data, a plurality of training successful trajectories, wherein each training successful trajectory is a successful trajectory that corresponds to a respective one of the training context inputs;

determining, using the auxiliary reward function and in accordance with current values of the auxiliary reward parameters, a respective auxiliary reward for each training successful trajectory;

training the policy to update current values of the policy parameters using the auxiliary rewards for the training successful trajectories;

obtaining, from the successful trajectory data, a plurality of validation successful trajectories, wherein each validation successful trajectory is a successful trajectory that corresponds to a respective one of the validation context inputs; and

training the auxiliary reward function on an objective that measures, for each validation successful trajectory, a probability assigned to the validation successful trajectory by outputs of the policy generated by processing the corresponding validation context input in accordance with the updated values of the policy parameters to update the current values of the auxiliary reward parameters, comprising computing gradients of the objective with respect to the auxiliary reward parameters using gradients of the objective with respect to the updated values of the policy parameters.

2 . The method of claim 1 , the following operations further comprising:

obtaining a training mini-batch of training context inputs and a validation mini-batch of validation context inputs;

for each context input in the training and validation mini-batches, generating a plurality of exploratory trajectories using the policy and in accordance with the current values of the policy parameters;

obtaining a respective task reward value for each of the exploratory trajectories that indicates whether the task was successfully completed by performing the exploratory trajectory; and

adding, to the successful trajectory data, any exploratory trajectories for which the task reward indicates that the task was successfully completed.

3 . The method of claim 1 , wherein the auxiliary reward function is a linear function that computes the auxiliary reward value as a linear combination of the task features based on the auxiliary reward parameters.

4 . The method of claim 1 , wherein the auxiliary reward function is a neural network that receives the task features as input and processes the task features in accordance with the auxiliary reward parameters to generate the auxiliary task reward.

5 . The method of claim 1 , wherein generating successful trajectory data that identifies a plurality of successful trajectories comprises:

training the policy starting from initial values of the policy parameters to generate first values of the policy parameters using only task rewards; and

periodically during the training using the policy to identify successful trajectories.

6 . The method of claim 5 , wherein training the policy starting from initial values of the policy parameters to generate first values of the policy parameters using only task rewards comprises training the policy to optimize an Iterative Maximum Likelihood (IML) objective.

7 . The method of claim 5 , wherein training the policy jointly with the auxiliary reward function comprises training the policy starting from the first values of the policy parameters.

8 . The method of claim 1 , wherein generating successful trajectory data that identifies a plurality of successful trajectories comprises:

identifying at least a subset of the trajectories in the successful trajectory data using random search.

9 . The method of claim 1 , wherein updating the current values of the policy parameters comprises updating the current values of the policy parameters to optimize a Memory Augmented Policy Optimization (MAPO) objective that is based on the auxiliary rewards.

10 . The method of claim 1 , wherein the objective measures, for each validation successful trajectory, a product of the received task reward value for the validation successful trajectory and a probability assigned to the validation successful trajectory by outputs of the policy generated by processing the corresponding validation context input in accordance with the updated values of the policy parameters to update the current values of the auxiliary reward parameters.

11 . A system for learning a policy having a plurality of policy parameters and configured to receive a context input for an agent to perform a task and to generate, from the context input and in accordance with the policy parameters, a trajectory of one or more actions for performing the task, wherein the one or more actions comprise control inputs for controlling the agent, the system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform operations comprising:

obtaining a plurality of training context inputs and a plurality of validation context inputs;

generating successful trajectory data that identifies a plurality of successful trajectories, each successful trajectory corresponding to a respective context input from either the training context inputs or the validation context inputs and being a trajectory for which a received task reward value indicates that the task was successfully completed by performing the trajectory in response to the corresponding context input; and

training the policy jointly with an auxiliary reward function having a plurality of auxiliary reward parameters, the auxiliary reward function being configured to map, in accordance with the auxiliary reward parameters, trajectory features of at least a trajectory to an auxiliary reward value that indicates how well the trajectory performed the task in response to the context input, and the training comprising repeatedly performing at least the following operations:

obtaining, from the successful trajectory data, a plurality of training successful trajectories, wherein each training successful trajectory is a successful trajectory that corresponds to a respective one of the training context inputs;

determining, using the auxiliary reward function and in accordance with current values of the auxiliary reward parameters, a respective auxiliary reward for each training successful trajectory;

training the policy to update current values of the policy parameters using the auxiliary rewards for the training successful trajectories;

obtaining, from the successful trajectory data, a plurality of validation successful trajectories, wherein each validation successful trajectory is a successful trajectory that corresponds to a respective one of the validation context inputs; and

training the auxiliary reward function on an objective that measures, for each validation successful trajectory, a probability assigned to the validation successful trajectory by outputs of the policy generated by processing the corresponding validation context input in accordance with the updated values of the policy parameters to update the current values of the auxiliary reward parameters, comprising computing gradients of the objective with respect to the auxiliary reward parameters using gradients of the objective with respect to the updated values of the policy parameters.

12 . The system of claim 11 , the following operations further comprising:

obtaining a training mini-batch of training context inputs and a validation mini-batch of validation context inputs;

for each context input in the training and validation mini-batches, generating a plurality of exploratory trajectories using the policy and in accordance with the current values of the policy parameters;

obtaining a respective task reward value for each of the exploratory trajectories that indicates whether the task was successfully completed by performing the exploratory trajectory; and

adding, to the successful trajectory data, any exploratory trajectories for which the task reward indicates that the task was successfully completed.

13 . The system of claim 11 , wherein the auxiliary reward function is a linear function that computes the auxiliary reward value as a linear combination of the task features based on the auxiliary reward parameters.

14 . The system of claim 11 , wherein the auxiliary reward function is a neural network that receives the task features as input and processes the task features in accordance with the auxiliary reward parameters to generate the auxiliary task reward.

15 . The system of claim 11 , wherein generating successful trajectory data that identifies a plurality of successful trajectories comprises:

training the policy starting from initial values of the policy parameters to generate first values of the policy parameters using only task rewards; and

periodically during the training using the policy to identify successful trajectories.

16 . The system of claim 15 , wherein training the policy starting from initial values of the policy parameters to generate first values of the policy parameters using only task rewards comprises training the policy to optimize an Iterative Maximum Likelihood (IML) objective.

17 . The system of claim 15 , wherein training the policy jointly with the auxiliary reward function comprises training the policy starting from the first values of the policy parameters.

18 . The system of claim 11 , wherein generating successful trajectory data that identifies a plurality of successful trajectories comprises:

identifying at least a subset of the trajectories in the successful trajectory data using random search.

19 . The system of claim 11 , wherein updating the current values of the policy parameters comprises updating the current values of the policy parameters to optimize a Memory Augmented Policy Optimization (MAPO) objective that is based on the auxiliary rewards.

20 . One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for learning a policy having a plurality of policy parameters and configured to receive a context input for an agent to perform a task and to generate, from the context input and in accordance with the policy parameters, a trajectory of one or more actions for performing the task, wherein the one or more actions comprise control inputs for controlling the agent, the operations comprising:

obtaining a plurality of training context inputs and a plurality of validation context inputs;

generating successful trajectory data that identifies a plurality of successful trajectories, each successful trajectory corresponding to a respective context input from either the training context inputs or the validation context inputs and being a trajectory for which a received task reward value indicates that the task was successfully completed by performing the trajectory in response to the corresponding context input; and

training the policy jointly with an auxiliary reward function having a plurality of auxiliary reward parameters, the auxiliary reward function being configured to map, in accordance with the auxiliary reward parameters, trajectory features of at least a trajectory to an auxiliary reward value that indicates how well the trajectory performed the task in response to the context input, and the training comprising repeatedly performing at least the following operations:

obtaining, from the successful trajectory data, a plurality of training successful trajectories, wherein each training successful trajectory is a successful trajectory that corresponds to a respective one of the training context inputs;

determining, using the auxiliary reward function and in accordance with current values of the auxiliary reward parameters, a respective auxiliary reward for each training successful trajectory;

training the policy to update current values of the policy parameters using the auxiliary rewards for the training successful trajectories;

obtaining, from the successful trajectory data, a plurality of validation successful trajectories, wherein each validation successful trajectory is a successful trajectory that corresponds to a respective one of the validation context inputs; and

training the auxiliary reward function on an objective that measures, for each validation successful trajectory, a probability assigned to the validation successful trajectory by outputs of the policy generated by processing the corresponding validation context input in accordance with the updated values of the policy parameters to update the current values of the auxiliary reward parameters, comprising computing gradients of the objective with respect to the auxiliary reward parameters using gradients of the objective with respect to the updated values of the policy parameters.