TRAINING MULTI-MODAL INTERACTIVE AGENTS USING A REWARD MODEL

1 . A method performed by one or more computers and for training a policy neural network for controlling an agent interacting with an environment, the method comprising:

at each of a plurality of time steps during a task episode:

obtaining an observation for the time step, the observation comprising:

an image characterizing a state of the environment at the time step; and

a natural language text sequence for the time step that characterizes a task being performed by the agent in the environment at the time step;

processing a policy input comprising the observation image and the natural language text sequence using the policy neural network to select one or more actions to be performed by the agent in response to the observation image; and

processing a reward input comprising the observation image and the natural language text sequence using a reward neural network, wherein the reward neural network is configured to process the observation image and the natural language text to generate a reward output that defines a reward value that characterizes a progress of the agent in performing the task characterized by the natural language text sequence as of the time step; and

training the policy neural network through reinforcement learning using the respective reward values for the plurality of time steps.

2 . The method of claim 1 , further comprising:

prior to training the policy neural network through reinforcement learning, training the policy neural network through imitation learning on data characterizing interactions between a plurality of agents in the environment.

3 . The method of claim 1 , wherein the reward neural network is configured to generate reward values that represent a utility to performing the task of a trajectory of observations up to and including the observation for the time step.

4 . The method of claim 1 , wherein training the policy neural network through reinforcement learning using the respective reward values for the plurality of time steps comprises:

for each of the plurality of time steps, computing a per-time step reward based on a difference between the reward value at the time step and the reward value for a preceding time step; and

training the policy neural network through reinforcement learning using the respective per-time step rewards for the plurality of time steps.

5 . The method of claim 1 , wherein the reward input further comprises a reward value generated for a preceding time step.

6 . The method of claim 1 , wherein the reward input further comprises a natural language output generated by the agent at a preceding time step.

7 . The method of claim 1 , wherein the reward output is the reward value.

8 . The method of claim 1 , wherein the reward output comprises a probability distribution over a set of possible reward values.

9 . The method of claim 1 , wherein the reward neural network has been trained on training data that includes, for each of one or more time steps in each of a plurality of training task episodes, a respective reward label selected from a set of reward labels that includes:

(i) a negative reward label that indicates that, as of the time step, the agent regressed from achieving a goal characterized in a natural language instruction for the time step, and

(ii) a positive reward label that indicates that, as of the time step, the agent made progress in achieving the goal characterized in the natural language instruction for the time step.

10 . The method of claim 9 , wherein the set also includes:

(iii) a neutral reward label.

11 . The method of claim 9 , wherein the reward neural network has been trained on the training data on a loss function that measures differences between reward values predicted for two time steps within the same training task episode.

12 . The method of claim 11 , wherein, for a given pair of time steps within the same training task episode:

if both time steps in the pair have a positive reward label and no time steps that are between the time steps in the pair have a negative reward label, the loss function encourages the reward value predicted for a later time step in the pair to be greater than the reward value predicted for an earlier time step in the pair; and

if both time steps in the pair have a negative reward label and no time steps that are between the time steps in the pair have a positive reward label, the loss function encourages the reward value predicted for the earlier time step in the pair to be greater than the reward value predicted for the later time step in the pair.

13 . The method of claim 12 , wherein, for the given pair of time steps:

if both time steps in the pair have a neutral reward label and all time steps that are between the time steps in the pair have a neutral reward label, the loss function encourages the reward value predicted for the later time step in the pair to be equal to the reward value predicted for the earlier time step in the pair.

14 . (canceled)

15 . The method of claim 9 , wherein processing a reward input comprising the observation image and the natural language text sequence using a reward neural network comprises:

using an image embedding neural network to generate a plurality of image embeddings that represent the observation image;

processing an input comprising the natural language text sequence using a text embedding neural network to generate a plurality of text embeddings that represent at least the natural language text sequence;

processing an input comprising the image embeddings and the text embeddings using a multi-modal neural network to generate an aggregated embedding; and

processing an input comprising the aggregated embedding using a reward neural network head to generate the reward value.

16 . The method of claim 15 , wherein the multi-modal neural network is a multi-modal Transformer neural network that is configured to (i) apply self-attention over at least the text embeddings and the image embeddings to generate respective updated embeddings for at least the plurality of text embeddings and (ii) generate the aggregated embedding from at least the respective updated embeddings for the text embeddings, wherein the input to the multi-modal Transformer neural network comprises the image embeddings, the text embeddings, and one or more dedicated embeddings.

17 . The method of claim 16 , wherein the multi-modal Transformer neural network comprises one or more self-attention layers that each have one or more self-attention heads, and wherein applying self-attention comprises processing the input through the one or more self-attention layers.

18 . (canceled)

19 . The method of claim 17 , wherein applying self-attention comprises generating respective updated embeddings for the text embeddings and the dedicated embeddings without updating the image embeddings.

20 . The method of claim 19 , wherein each self-attention head of each self-attention layer is configured to:

receive a head input comprising (i) the image embeddings generated by the image embedding neural network and (ii) respective current embeddings for the text embeddings and the dedicated embeddings;

generate, from the respective current embeddings, a respective query corresponding to each text embedding and each dedicated embedding;

generate, from the image embeddings and the respective current embeddings, a respective key corresponding to each image embedding, each text embedding, and each dedicated embedding;

generate, from the image embeddings and the respective current embeddings, a respective value corresponding to each image embedding, each text embedding, and each dedicated embedding; and

apply query-key-value attention over the respective queries, keys, and values to generate a respective initial updated embedding for each text embedding and each dedicated embedding without updating the image embeddings.

21 . The method of claim 20 , wherein generating the aggregated embedding comprises:

aggregating the respective updated embeddings for the text embeddings and the dedicated embeddings to generate an initial aggregated embedding; and

combining the respective updated embeddings for the dedicated embeddings with the initial aggregated embedding to generate the aggregated embedding, wherein the combining comprises concatenating each respective updated embedding and the initial aggregated embedding.

22 . (canceled)

23 . The method of claim 15 , wherein processing an input comprising the aggregated embedding using a reward neural network head to generate the reward output comprises:

generating a state representation from the aggregated embedding, the generating comprising processing the aggregated embedding using a memory neural network, wherein the memory neural network is a recurrent neural network; and

processing the state representation using one or more neural network layers to generate the reward output.

24 . (canceled)

25 . (canceled)

26 . The method of claim 15 , wherein the reward input further comprises a natural language output generated by the agent at a preceding time step, wherein the input to the text embedding neural network further comprises the natural language output and the text embeddings represent the natural language text sequence and the natural language output.

27 . The method of claim 1 , wherein the natural language text sequence is generated from a natural language text sequence that is generated based on a corresponding natural language text sequence from a corresponding time step in a training task episode.

28 . The method of claim 1 , wherein the natural language text sequence is generated by a setter agent in the environment, the setter agent controlled using a setter neural network, wherein the setter neural network has been trained to imitate an expert setter agent through imitation learning.

29 . (canceled)

30 . (canceled)

31 . The method of claim 15 , wherein the reward neural network has been trained on an overall loss function that includes (i) a loss function that is based on the reward labels and (ii) one or more auxiliary losses comprising an imitation learning loss that is computed using an output generated by an auxiliary policy neural network head that generates policy outputs for controlling the agent and/or a contrastive self-supervised representation learning loss.

32 . (canceled)

33 . (canceled)

34 . The method of claim 31 , wherein the one or more auxiliary losses include a cross-modality matching loss that uses outputs generated by the multi-modal neural network.

35 . The method of claim 2 , wherein at least some parameter values of the reward neural network were initialized using parameter values of the policy neural network that were determined by training the policy neural network through imitation learning.

36 . The method of claim 1 , wherein training the policy neural network through reinforcement learning using the respective reward values for the plurality of time steps comprises:

training the policy neural network through reinforcement learning and through imitation learning on an imitation learning data set.

37 . (canceled)

38 . (canceled)

39 . (canceled)

40 . The method of claim 1 , wherein the environment is a computing environment and the agent is a software agent executing within the computing environment to control, by performing actions selected by the trained policy neural network, one or more computing devices to carry out a task specified by a user interacting with the software agent.

41 . A 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 more computers to perform the operations, the operations comprising:

at each of a plurality of time steps during a task episode:

obtaining an observation for the time step, the observation comprising:

an image characterizing a state of the environment at the time step; and

a natural language text sequence for the time step that characterizes a task being performed by the agent in the environment at the time step;

processing a policy input comprising the observation image and the natural language text sequence using the policy neural network to select one or more actions to be performed by the agent in response to the observation image; and

processing a reward input comprising the observation image and the natural language text sequence using a reward neural network, wherein the reward neural network is configured to process the observation image and the natural language text to generate a reward output that defines a reward value that characterizes a progress of the agent in performing the task characterized by the natural language text sequence as of the time step; and

training the policy neural network through reinforcement learning using the respective reward values for the plurality of time steps.

42 . One or more computer storage media storing instructions that when executed by one or more computers cause the one more computers to perform operations, the operations comprising:

at each of a plurality of time steps during a task episode:

obtaining an observation for the time step, the observation comprising:

an image characterizing a state of the environment at the time step; and

a natural language text sequence for the time step that characterizes a task being performed by the agent in the environment at the time step;

processing a policy input comprising the observation image and the natural language text sequence using the policy neural network to select one or more actions to be performed by the agent in response to the observation image; and

processing a reward input comprising the observation image and the natural language text sequence using a reward neural network, wherein the reward neural network is configured to process the observation image and the natural language text to generate a reward output that defines a reward value that characterizes a progress of the agent in performing the task characterized by the natural language text sequence as of the time step; and

training the policy neural network through reinforcement learning using the respective reward values for the plurality of time steps.