A Dual-Critic Reinforcement Learning Framework for Frame-level Bit Allocation in HEVC/H.265

The frame-level bit allocation is to allocate a proper number of bits to every video frame by choosing a QP for its encoding, in order to minimize the distortion of a GOP subject to a rate constraint. In symbols, we have

where $QP_{i}$ denotes the QP chosen for the $i$-th frame in a GOP, $D_{i}(QP_{i})$ is the distortion incurred by encoding frame $i$ with $QP_{i}$, $R_{i}(QP_{i})$ is the actual number of bits produced by the encoder, and $R_{\text{GOP}}$ is the GOP-level bit budget.

In this paper, we tackle the learning of an agent that can perform bit allocation among video frames in a GOP from an RL perspective. The frame-level bit allocation for a GOP is regarded as an episodic task, with the learning process given as follows and illustrated in Fig. 1. (1) First, we evaluate a state signal (Section 2.2) for an input frame. (2) Second, the state signal is fed to the agent implemented by a neural network (Section 2.5) to determine a QP for encoding the current frame. (3) Third, based on the chosen QP, the current frame is encoded and decoded by the specified codec (e.g. x265) to update the state signal for the agent to decide the QP for the next video frame. Moreover, two different types of rewards (i.e. the distortion and the rate rewards) are computed to indicate how good the choice of QP is (Section 2.3). These steps are repeated iteratively until a terminal state (i.e. the end of an episode) is reached. The training has the agent interact with the codec many times so that it can learn to maximize the rewards (Section 2.4).

The state signal serves as the input to the agent for its decision making. We follow mainly the state design in [1], which is summarized in Table 1 for easy reference. It involves several hand-crafted intra- and inter-frame features that characterize the statistics of all the video frames in a GOP and their residual frames produced with zero motion vectors for simplicity. The fact that the state depends on the statistics of the remaining video frames in a GOP suggests a look-ahead strategy. As such, when a reference frame is not yet coded, we turn to its original and uncompressed version to compute the residual frame.

The reward signal plays a central role in shaping the agent’s behavior. The agent is usually trained to maximize a single reward. Because our task is to minimize the distortion of a GOP subject to a rate constraint, we develop two rewards: the distortion $r_{D}$ and the rate $r_{R}$ rewards. The former $r_{D}$ is evaluated upon the completion of encoding a video frame as an immediate reward, reflecting how good the current choice of QP is. It is defined as the normalized negative distortion of a compressed video frame in mean square error (MSE):

where $MSE_{i}$ is the MSE of the current frame $i$, and the $MSE_{QP0}$ and $MSE_{QP51}$ are the MSE’s of the entire GOP by encoding every video frame with the two extreme QP values, i.e. QP0 and QP51. The two normalization factors cause the sum of $r_{D_{i}}$ over the video frames in a GOP to fall approximately in the interval from 0 to 1. Unlike the distortion reward, which is evaluated for every video frame, the rate reward $r_{R}$ is given at the end of an episode, capturing the absolute bias between the actual bit rate and the target bit rate. It is specified by

To train the proposed agent following the aforementioned state and reward design, we depart from the single critic approach to learn two critics, with one $Q_{D}$ used for predicting the distortion reward (more precisely, the distortion-to-go $\sum_{i}r_{D_{i}}$) and the other $Q_{R}$ for predicting the rate reward $r_{R}$. The agent is updated adaptively with one of the critics. Specifically, we constantly apply the agent under training to encode a GOP. If it produces a bit rate exceeding the target $R_{GOP}$, we use the rate critic $Q_{R}$ to update its network parameters, with the aim of optimizing the agent towards precise rate control. In the other case, where the actual bit rate is lower than the target, it is updated with the distortion critic $Q_{D}$ to minimize the distortion of the reconstructed GOP.

Algorithm 1 details the training process. It can be divided into two major parts: Part I (lines 5 to 18) and Part II (lines 19 to 30). Part I corresponds to a rollout of the policy together with a noise process to collect transitions of states, actions, and rewards in a replay buffer $R_{C}$, the data of which are utilized to update the two critics $Q_{D}$ and $Q_{R}$ based on the ordinary DDPG algorithm. Note that a zero immediate rate reward is recorded at every step until the very last step (the terminal step/state), at which the bit rate bias $r_{R}$ of the entire GOP is given as the immediate rate reward. Part II is another rollout of the policy but without the noise process. That is, the agent is put to use as for inference. The state transitions experienced is stored in a separate replay buffer $R_{A}$. At the end of the rollout, the actual bit rate used is compared against the target bit rate to decide which of the two critics should be used to update the agent.

As Fig. 1 illustrates, our agent consists of two modules: the base and the actor modules. The base module is pre-trained to mimic the QP control of the selected codec (e.g. x265), while the actor module is learned with our dual-critic RL framework to update the QP of the base module by a delta QP in the interval $[-10,10]$. In training the actor module, the base module is freezed. The introduction of the base module restricts the exploration space for learning the actor module. This is found beneficial particularly because the convergence can be an issue for dual-critic training with a vast exploration space of QP options. The architectures of our actor and critic networks are detailed in Table 2. Note that the base module shares the same architecture as the actor module, except for the exclusion of the base QP as input (Fig. 1).

Our scheme differs from the single-critic design [1] in two aspects. First, it learns two separate critics, whereas Chen et al. [1] use one critic to learn a single reward function for a GOP as the weighted combination of the distortion and rate rewards, namely $-1\times\sum_{i=1}^{N}r_{D_{i}}-\lambda\times r_{R}$, where $\lambda$ is a hyper-parameter chosen empirically to trade the distortion of a GOP against the rate deviation from the target $R_{GOP}$. We argue that the use of a fixed $\lambda$, as adopted in [1], could hardly learn an agent that works well on different types of videos. Second, we require our agent to determine QP’s for the remaining frames when the target $R_{GOP}$ is exceeded, whereas Chen et al. [1] propose encoding the remaining frames in terminal mode with $QP=QP_{I}+10$, with $QP_{I}$ denoting the I-frame’s QP.

We assess the objective and subjective compression performance of the proposed method, with the results compared against those produced by x265 [7], the base module alone (referred hereafter to as base), and the single-critic method [1]. The baseline methods do not include [3] because it is optimized specifically for low-delay p-frame coding, as compared to our GOP coding with hierarchical bi-prediction.

All the competing methods are evaluated on x265, with a hierarchical GOP coding structure as depicted in Fig. 2. Because our focus is on the bit allocation inside a GOP, all the tested methods follow the same GOP-level bit allocation as x265. This is achieved by first encoding every test sequence with fixed QP 22, 27, 32, and 37 to establish four sequence-level target bit rates $R_{s}$. They are then utilized as the sequence-level rate constraints to encode every test sequence again by turning on the ABR rate control of x265 (–bitrate $R_{s}$ –vbv-bufsize 2*$R_{s}$ –vbv-maxrate $R_{s}$), the results of which serve as our x265 baseline. For a fair comparison, the GOP-level bit rates produced by x265 in this ABR mode are taken as the GOP rate constraint $R_{GOP}$ for the other methods and the encoding is restricted to one pass only.

We train separate models for the four target bit rates. The base module clones approximately the rate control behavior of x265 using its QP choices as the supervision ground truths. For a fair comparison, both our scheme and the single-critic method  [1] operate with the same base module. Our training datasets include UVG [8], MCL-JCV [9], and Class A sequences in JCT-VC dataset. At test time, we use Class B and Class C sequences in JCT-VC dataset, which are not seen at training time. To speed up the RL training processes, the training and test videos are downscaled to $512\times 320$. The run time of our feature extraction and actor network forwarding only occupies 4% run time of ABR rate control of x265.

The objective compression performance is reported in terms of BD-PSNR gains and BD-rate savings, with x265 in ABR mode serving as anchor. The quality metrics are Y-PSNR and YUV-PSNR evaluated (and averaged) over individual frames, where YUV-PSNR is given by $(6\times Y\_PSNR+U\_PSNR+V\_PSNR)/8$. The accuracy in rate control is quantified by averaging the absolute rate deviations from $R_{GOP}$ in percentage terms over all the GOPs in a test sequence.

Fig. 3 compares the rate-distortion (R-D) curves of different methods for three selected sequences, with the complete BD-rate and BD-PSNR results summarized in Table 3. Also compared in the same table is the rate control accuracy averaged over all the rate points. From these results, the following observations are immediate:

(1) Our method outperforms the competing methods in terms of R-D performance. Fig. 3 shows that the proposed method achieves superior R-D performance to x265, as is confirmed by its BD-PSNR and BD-rate numbers in Table 3. It improves Y-PSNR by 0.75dB and YUV-PSNR by 0.91dB, as compared to 0.37dB and 0.39dB improvements by Chen et al. [1], respectively. The corresponding BD-rate savings range from 14.7% to 18.8%, which nearly double the rate savings by Chen et al. [1] (7.6% to 8.8%).

(2) The single-critic method has inconsistent R-D performance. From Fig. 3, it performs better than x265 in slow-motion sequences (e.g. BQTerrace), but worse in the other fast-motion sequences (e.g. BasketballDrill, RaceHorses). The performance degradation, if any, is more obvious at higher bit rates. Its inconsistent behavior in different types of videos and at different bit rates may be attributed to the use of a fixed hyper-parameter (Section 2.6) for trading off the video quality against the rate penalty, which is unable to generalize well to all the scenarios.

(3) The base module is unable to meet the GOP-level rate constraints although showing comparable or better R-D performance than x265. Table 3 reveals that its average deviation from the GOP-level bit rate constraint is around 14% and can be as high as nearly 25% in some sequences. However, with the RL training, both the single-critic method [1] and ours can match the GOP-level bit rates of x265 more closely (the left plot of Fig. 4), showing much reduced rate deviations of 6.8% and 6.2% (Table 3), respectively. From Table 3, the base module achieves 4.9-5.0% BD-rate savings and 0.22-0.24dB BD-PSNR gains over x265. This is because the ABR rate control may suffer from an unstable transient behavior in encoding the first few GOP’s, as indicated by the developers of x265 [7]. By excluding these GOP’s from behavior cloning, the base module learns a somewhat fixed and regular QP pattern over different GOP’s, as shown in the right plot of Fig. 4. Nevertheless, the higher rate control accuracy and much improved R-D performance of the proposed method than the base module stress the contributions of our RL framework (Table 3).

Fig. 5 visualizes the frame-level QP assignment within GOP’s along with the corresponding coded frame size and YUV-PSNR. We make the following observations:

(4) Our RL agent tends to choose smaller QP values for I- and B-frames, allocating more bits to their coding. In slow-motion sequences (e.g. BQTerrace), our approach encodes I-frames at an even higher quality using an extremely small QP. Such a policy is intuitively agreeable since I- and B-frames serve as reference frames for the remaining non-reference b-frames. In slow-motion sequences, the quality of I-frames is even more critical to the reconstruction quality of a GOP.

(5) x265 and its clone version, i.e. the base module, focus more on I-frames without showing a drastic difference in QP selection for B- and b-frames. A similar policy is observed in both fast- and slow-motion sequences. This causes B- and b-frames to have similar YUV-PSNR, and B-frames to use slightly more bits in fast-motion sequences (e.g. RaceHorses) due to a longer-distance prediction.

(6) The single-critic method [1] learns a subtle policy on QP selection. In slow-motion sequences (e.g. BQTerrace), it behaves similarly to our scheme except that it does not particularly favor B-frames and enters terminal state early by encoding the remaining frames with $QP=QP_{I}+10$ (Section 2.6). Interestingly, such a policy shows comparable R-D performance to ours (BQTerrace in Table 3). However, in fast-motion sequences, I-frames are coded with relatively larger QP’s than B- and b-frames, which explains its poorer R-D performance in these sequences.

(7) Our RL agent treats b-frames differently in fast-motion sequences. It shows uneven QP assignment among b-frames, with their QP’s increasing with the coding order (Fig. 2). Recall that our objective is to minimize the GOP-level distortion; there is no constraint on the frame-level quality distribution in a GOP. From Fig. 2, b-frames are coded independently of each other as non-reference frames. The QP assignment in Fig. 4 turns out to be a reasonably good solution since the GOP-level distortion is found to be smaller than those of the other competing methods.

Fig. 6 presents randomly chosen sample images produced by the competing methods. We observe that the proposed method retains more texture details, e.g. the wooden stripe in BasketballDrill, the crease of the shirt in BQMall, and the grass region in RaceHorses. Videos are provided at http://mapl.nctu.edu.tw/RL_Rate_Control/.