Gated linear contextual bandits
Methods, systems, and apparatus, including computer programs encoded on computer-readable storage media, for training a neural network to control a real-world agent interacting with a real-world environment to cause the real-world agent to perform a particular task. One of the methods includes training the neural network to determine first values of the parameters by optimizing a first task-specific objective that measures a performance of the policy neural network in controlling a simulated version of the real-world agent; obtaining real-world data generated from interactions of the real-world agent with the real-world environment; and training the neural network to determine trained values of the parameters from the first values of the parameters by jointly optimizing (i) a self-supervised objective that measures at least a performance of internal representations generated by the neural network on a self-supervised task performed on the real-world data and (ii) a second task-specific objective.
1 . A method of selecting an action from a set of actions to be performed in response to each context in a sequence of contexts, the method comprising:
maintaining data specifying a respective gated linear network corresponding to each action of the set of actions, wherein each gated linear network is configured to predict a probability that a reward will be received if the corresponding action is performed in response to an input context, and wherein each gated linear network comprises a plurality of layers, each layer comprising one or more neurons, wherein each neuron in each layer after the first layer is configured to receive (i) the input context and (ii) predictions from neurons in a previous layer and to (iii) apply a gating function to the input context to select a weight vector and (iv) generate as output a geometric mixture of the predictions from the neurons in the previous layer based on the selected weight vector;
for each context in the sequence of contexts:
for each action, processing the context using the gated linear network corresponding to the action to generate a predicted probability;
for each action, generating an action score for the action from at least the predicted probability, comprising:
computing a pseudo-count for the action, wherein computing the pseudo-count comprises determining an overlap between (i) a signature of the context across the gating functions of the neurons in the gated linear network for the action and (ii) signatures of any earlier contexts in the sequence for which the action was selected as the action to be performed in response to the earlier context, wherein the signature of the context characterizes outputs of the gating functions of the neurons in the gated linear network corresponding to the action generated by processing the context; and
generating the action score from the predicted probability for the action and the pseudo-count for the action; and
selecting the action to be performed in response to the context based on the action scores.
2 . The method of claim 1 , wherein selecting the action to be performed in response to the context based on the action scores comprises selecting the action with a highest action score.
3 . The method of claim 1 , further comprising:
for each context in the sequence of contexts:
receiving a reward; and
updating the gated linear network for the selected action based on the reward.
4 . The method of claim 3 , wherein updating the gated linear network for the selected action comprises:
updating each neuron in the gated linear network locally based on a neuron-specific loss.
5 . The method of claim 1 , wherein a last layer of the plurality of layers includes only a single neuron, and wherein the predicted probability of the gated linear network is the output of the single neuron.
6 . The method of claim 1 , wherein the neurons in the first layer of the plurality of layers receive the input context and a set of base predictions.
7 . The method of claim 1 , wherein determining the overlap comprises:
determining a count of how many earlier contexts satisfy both (1) the action was selected in response to the earlier context and (2) for each neuron in the gated linear network for the action, the output of the gating function for the neuron for the earlier context is the same as the output of the gating function for the neuron for a current context.
8 . The method of claim 1 , wherein the set of actions are recommendations of content items.
9 . A method of selecting an action from a set of actions to be performed in response to each context in a sequence of contexts, the method comprising:
maintaining data specifying a respective tree of gated linear networks corresponding to each action of the set of actions, wherein each tree of gated linear networks is collectively configured to predict a respective probability for each of a plurality of bins of a range of reward values, wherein the respective probability for each bin represents a likelihood that a reward that falls in the bin will be received if the corresponding action is performed in response to an input context, and wherein each gated linear network comprises a plurality of layers, each layer comprising one or more neurons, wherein each neuron in each layer after the first layer is configured to receive (i) the input context and (ii) predictions from neurons in a previous layer and to (iii) apply a gating function to the input context to select a weight vector and (iv) generate as output a geometric mixture of the predictions from the neurons in the previous layer based on the selected weight vector;
for each context in the sequence of contexts:
for each action, processing the context using the tree of gated linear networks corresponding to the action to generate a respective probability for each of the plurality of bins of the range of reward values;
for each action, generating an action score for the action from at least the respective probabilities, comprising:
computing a pseudo-count for the action, the computing the pseudo-count comprises determining an overlap between (i) a signature of the context across the gating functions of the neurons in the gated linear networks in the tree of gated linear networks for the action and (ii) signatures of the context across the gating functions of the neurons in the gated linear networks in the tree of gated linear networks for any earlier contexts in the sequence for which the action was selected as the action to be performed in response to the earlier context, wherein the signature of a context refers to the outputs of the gating functions of the neurons in the gated linear networks in the tree of gated linear networks for the action generated by processing the context; and
selecting the action to be performed in response to the context based on the action scores.
10 . The method of claim 9 , wherein selecting the action to be performed in response to the context based on the action scores comprises selecting the action with a highest action score.
11 . The method of claim 9 , further comprising: for each context in the sequence of contexts:
receiving a reward; and
updating the tree of gated linear networks for the selected action based on the reward.
12 . The method of claim 11 , wherein updating the tree of gated linear networks for the selected action comprises:
for each neuron on a path to an identified bin through the tree of gated linear networks for the selected action, updating the neuron locally based on a neuron-specific loss.
13 . The method of claim 9 , wherein a last layer of the plurality of layers includes only a single neuron, and wherein the predicted probability of the gated linear network is the output of the single neuron.
14 . The method of claim 9 , wherein the neurons in the first layer of the plurality of layers receive the input context and a set of base predictions.
15 . A system comprising one or more computers and one or more storage devices storing instructions that when executed by one or more computers cause the one or more computers to perform operations for selecting an action from a set of actions to be performed in response to each context in a sequence of contexts, the operations comprising:
maintaining data specifying a respective gated linear network corresponding to each action of the set of actions, wherein each gated linear network is configured to predict a probability that a reward will be received if the corresponding action is performed in response to an input context, and wherein each gated linear network comprises a plurality of layers, each layer comprising one or more neurons, wherein each neuron in each layer after the first layer is configured to receive (i) the input context and (ii) predictions from neurons in a previous layer and to (iii) apply a gating function to the input context to select a weight vector and (iv) generate as output a geometric mixture of the predictions from the neurons in the previous layer based on the selected weight vector;
for each context in the sequence of contexts:
for each action, processing the context using the gated linear network corresponding to the action to generate a predicted probability;
for each action, generating an action score for the action from at least the predicted probability, comprising:
computing a pseudo-count for the action, wherein computing the pseudo-count comprises determining an overlap between (i) a signature of the context across the gating functions of the neurons in the gated linear network for the action and (ii) signatures of any earlier contexts in the sequence for which the action was selected as the action to be performed in response to the earlier context, wherein the signature of the context characterizes outputs of the gating functions of the neurons in the gated linear network corresponding to the action generated by processing the context; and
generating the action score from the predicted probability for the action and the pseudo-count for the action; and
selecting the action to be performed in response to the context based on the action scores.
16 . The system of claim 15 , wherein selecting the action to be performed in response to the context based on the action scores comprises selecting the action with a highest action score.
17 . The system of claim 15 , the operations further comprising:
for each context in the sequence of contexts:
receiving a reward; and
updating the gated linear network for the selected action based on the reward.
18 . The system of claim 17 , wherein updating the gated linear network for the selected action comprises:
updating each neuron in the gated linear network locally based on a neuron-specific loss.
19 . The system of claim 15 , wherein a last layer of the plurality of layers includes only a single neuron, and wherein the predicted probability of the gated linear network is the output of the single neuron.
20 . The system of claim 15 , wherein the neurons in the first layer of the plurality of layers receive the input context and a set of base predictions.