Training rate control neural networks through reinforcement learning
Systems and methods for training rate control neural networks through reinforcement learning. During training, reward values for training examples are generated from the current performance of the rate control neural network in encoding the video in the training example and the historical performance of the rate control neural network in encoding the video in the training example.
1 . A method of training a rate control neural network for use in encoding sequences of video frames, wherein the rate control neural network is configured to, during encoding of an input sequence of video frames and for each video frame in the sequence, receive input features for the video frame and process the input features to generate a respective score for each of a plurality of possible quantization parameter values, the method comprising, at an iteration which is one of a sequence of iterations of training the rate control neural network:
obtaining a training example comprising a training sequence of video frames and a target data measure for the training sequence;
obtaining a historical performance for the training example, wherein the historical performance is a historical performance of training the rate control neural network on the training example at preceding iterations of the sequence of iterations;
generating an encoded representation of the training sequence that comprises a sequence of encoded video frames, comprising, for each video frame in the training sequence:
generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values;
selecting, using the policy output, a quantization parameter value from the plurality of possible quantization parameter values; and
encoding the video frame using a video codec in accordance with the selected quantization parameter value;
determining a current performance for the training example, comprising determining (i) a quality measure that measures a quality of the encoded representation and (ii) a data measure that measures an amount of data required to represent the encoded representation;
generating a reward value from the current performance and the historical performance; and
training the rate control neural network using the reward value through reinforcement learning.
2 . The method of claim 1 , wherein generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values, comprises:
processing input features for the video frame using the rate control neural network to generate a respective score for each of the plurality of possible quantization parameter values; and
using the respective scores as the respective policy scores.
3 . The method of claim 1 , wherein generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values, comprises:
performing a look-ahead search of possible future states of the encoding guided by the rate control neural network.
4 . The method of claim 1 , wherein the historical performance for the training example is generated from (i) an exponential moving average of previous quality measures computed for the training example during the training of the rate control neural network and (ii) an exponential moving average of previous overshoots determined from the target data measure and previous data measures computed for the training example during the training of the rate control neural network.
5 . The method of claim 4 , further comprising:
updating the historical performance using (i) the quality measure that measures a quality of the encoded representation and (ii) an overshoot for the data measure that measures a difference between an amount of data required to represent the encoded representation and the target data measure.
6 . The method of claim 4 , wherein generating a reward value from the current performance and the historical performance comprises:
determining whether either a current overshoot or the exponential moving average of previous overshoots are greater than zero;
when either a current overshoot or the exponential moving average of previous overshoots are greater than zero:
setting the reward value to a positive value if the current overshoot is less than or equal to the exponential moving average of previous overshoots, and
setting the reward value to a negative value if the current overshoot is greater than the exponential moving average of previous overshoots.
7 . The method of claim 6 , wherein generating a reward value from the current performance and the historical performance comprises:
when neither the current overshoot nor exponential moving average of previous overshoots are greater than zero:
setting the reward equal to the positive value if the quality measure that measures a quality of the encoded representation is the same as or better than the exponential moving average of previous quality measures, and
setting the reward value to the negative value if the quality measure that measures a quality of the encoded representation is worse than the exponential moving average of previous quality measures.
8 . The method of claim 1 , wherein generating a reward value from the current performance and the historical performance comprises:
setting the reward value to a positive value when a current performance measure is greater than a historical performance measure by more than a first threshold amount.
9 . The method of claim 8 , wherein generating a reward value from the current performance and the historical performance comprises:
setting the reward value to a negative value when the current performance measure is less than the historical performance measure by more than the first threshold amount.
10 . The method of claim 1 , wherein the rate control neural network comprises:
a representation neural network that is configured to process the input features to generate a feature representation for the video frame; and
a policy neural network head that is configured to process the feature representation to generate the respective score for each of the plurality of possible quantization parameter values.
11 . The method of claim 10 , further comprising, for each of one or more auxiliary tasks:
processing the feature representation for the video frame using a corresponding auxiliary neural network head to generate a predicted output for the auxiliary task; and
determining a respective error between the predicted output and a corresponding ground truth output for the auxiliary task;
wherein training the rate control neural network using the reward value through reinforcement learning comprises training the rate control neural network using the respective errors for the one or more auxiliary tasks.
12 . 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 or more computers to perform operations for training a rate control neural network for use in encoding sequences of video frames, wherein the rate control neural network is configured to, during encoding of an input sequence of video frames and for each video frame in the sequence, receive input features for the video frame and process the input features to generate a respective score for each of a plurality of possible quantization parameter values, the operations comprising, at an iteration which is one of a sequence of iterations of training the rate control neural network:
obtaining a training example comprising a training sequence of video frames and a target data measure for the training sequence;
obtaining a historical performance for the training example, wherein the historical performance is a historical performance of training the rate control neural network on the training example at preceding iterations of the sequence of iterations;
generating an encoded representation of the training sequence that comprises a sequence of encoded video frames, comprising, for each video frame in the training sequence:
generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values;
selecting, using the policy output, a quantization parameter value from the plurality of possible quantization parameter values; and
encoding the video frame using a video codec in accordance with the selected quantization parameter value;
determining a current performance for the training example, comprising determining (i) a quality measure that measures a quality of the encoded representation and (ii) a data measure that measures an amount of data required to represent the encoded representation;
generating a reward value from the current performance and the historical performance; and
training the rate control neural network using the reward value through reinforcement learning.
13 . The system of claim 12 , wherein generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values, comprises:
processing input features for the video frame using the rate control neural network to generate a respective score for each of the plurality of possible quantization parameter values; and
using the respective scores as the respective policy scores.
14 . The system of claim 12 , wherein generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values, comprises:
performing a look-ahead search of possible future states of the encoding guided by the rate control neural network.
15 . The system of claim 12 , wherein the historical performance for the training example is generated from (i) an exponential moving average of previous quality measures computed for the training example during the training of the rate control neural network and (ii) an exponential moving average of previous overshoots determined from the target data measure and previous data measures computed for the training example during the training of the rate control neural network.
16 . The system of claim 15 , the operations further comprising:
updating the historical performance using (i) the quality measure that measures a quality of the encoded representation and (ii) an overshoot for the data measure that measures a difference between an amount of data required to represent the encoded representation and the target data measure.
17 . The system of claim 12 , the operations further comprising, for each of one or more auxiliary tasks:
processing the feature representation for the video frame using a corresponding auxiliary neural network head to generate a predicted output for the auxiliary task; and
determining a respective error between the predicted output and a corresponding ground truth output for the auxiliary task;
wherein training the rate control neural network using the reward value through reinforcement learning comprises training the rate control neural network using the respective errors for the one or more auxiliary tasks.
18 . One or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for training a rate control neural network for use in encoding sequences of video frames, wherein the rate control neural network is configured to, during encoding of an input sequence of video frames and for each video frame in the sequence, receive input features for the video frame and process the input features to generate a respective score for each of a plurality of possible quantization parameter values, the method comprising, at an iteration which is one of a sequence of iterations of training the rate control neural network:
obtaining a training example comprising a training sequence of video frames and a target data measure for the training sequence;
obtaining a historical performance for the training example, wherein the historical performance is a historical performance of training the rate control neural network on the training example at preceding iterations of the sequence of iterations;
generating an encoded representation of the training sequence that comprises a sequence of encoded video frames, comprising, for each video frame in the training sequence:
generating, using the rate control neural network, a policy output for the video frame that comprises a respective policy score for each of the plurality of possible quantization parameter values;
selecting, using the policy output, a quantization parameter value from the plurality of possible quantization parameter values; and
encoding the video frame using a video codec in accordance with the selected quantization parameter value;
determining a current performance for the training example, comprising determining (i) a quality measure that measures a quality of the encoded representation and (ii) a data measure that measures an amount of data required to represent the encoded representation;
generating a reward value from the current performance and the historical performance; and
training the rate control neural network using the reward value through reinforcement learning.