Action-Constrained Reinforcement Learning for Frame-Level Bit Allocation in HEVC/H.265 through Frank-Wolfe Policy Optimization

The objective of the frame-level bit allocation is to minimize the distortion (i.e. maximize the VMAF score) of a GOP subject to a rate constraint. This is achieved by allocating available bits properly among video frames in a GOP through choosing their quantization parameters (QP). In symbols, we have

where $QP_{i}$ indicates the QP for the i-th frame, $N$ denotes the GOP size, $D_{i}(QP_{i})$ is the distortion (measured in VMAF) of frame $i$ encoded with $QP_{i}$, $R_{i}(QP_{i})$ is the number of encoded bits of frame $i$, and $R_{GOP}$ is the GOP-level rate constraint.

We address the problem in Eq. (6) by learning an RL agent that is able to determine in sequential steps the QP of every video frame in a GOP, so that the resulting number of encoding bits approximates $R_{GOP}$ while maximizing the cumulative VMAF score. To this end, we introduce the NFWPO-based RL framework.

We begin by drawing an analogy between the NFWPO-based RL setup and the problem in Eq. (6). Consider a video frame $i$ whose $QP_{i}$ is to be determined. First, the minimization of the cumulative distortion over the entire GOP in Eq. (6) is regarded as the maximization of the reward-to-go $Q$ in Eq. (1) in determining $QP_{i}$, where $Q$ in our problem corresponds to the sum of VMAF scores from frame $i$ till the very last frame $N$ in the GOP. That is, the local QP decision for frame $i$ needs to ensure that the reward in the long run (respectively, the cumulative distortion till the very last frame) can be maximized (respectively, minimized). In addition, in determining $QP_{i}$, we wish to satisfy the rate constraint in Eq. (6). Obviously, not every possible $QP_{i}$ can meet the constraint. Therefore, $QP_{i}$ needs to be confined in a feasible set $\mathcal{C}(s_{i})$, which is a function of the current coding context represented by a state signal $s_{i}$. How $\mathcal{C}(s_{i})$ is specified is detailed in Section 3.6. With these in mind, we translate the problem in Eq. (6) into

which shares the same form as Eq. (1). We can thus learn an actor network for determining $QP_{i}$ based on the action-constrained RL setup (Section 2).

Fig. 1 presents an overview of our action-constrained RL setup. When encoding frame $i$, (1) a state $s_{i}$ is first evaluated (Section 3.3). (2) It is then taken by the neural network-based RL agent to output $QP_{i}$ as an action (Section 3.4). (3) With $QP_{i}$, frame $i$ is encoded and decoded by the specified codec (e.g. x265). Upon the completion of encoding frame $i$, our proposed NFWPO-based RL algorithm evaluates a distortion reward $r_{D_{i+1}}$ and a rate reward $r_{R_{i+1}}$ (Section 3.5). These steps are repeated iteratively until all the video frames in a GOP are encoded.

At training time, the agent interacts with the codec intensively by encoding every GOP as an episodic task. To predict the distortion and rate rewards, we resort to two neural network-based critics, the distortion critic $Q_{D}(s,QP)$ and the rate critic $Q_{R}(s,QP)$, both of which take the state signal $s$ and $QP$ as inputs. As will be seen, the rate critic $Q_{R}(s,QP)$, which predicts the rate reward-to-go (or the rate deviation from the bit budget $R_{GOP}$ at the end of encoding a GOP), allows us to specify a feasible set $\mathcal{C}(s_{i})$ of $QP$ for the current coding frame. With $\mathcal{C}(s_{i})$, we apply NFWPO [6] to train the agent, in order to maximize the distortion reward-to-go $Q_{D}(s,QP)$, which gives an estimate of the cumulative VMAF score. That is, we substitute $Q_{D}$ for $Q$ in Eq. (7).

The state signal provides informative information for the agent to make decisions. Inspired by [1], we form our state signal with the following hand-crafted features. (1) The intra-frame feature calculates the variance of pixel values for the current frame. Likewise, (2) the inter-frame feature evaluates the mean of pixel values for the frame differences based on uni- and bi-prediction, where we adopt zero motion compensation as a compromise between performance and complexity. (3) The average of the intra-frame features averages the intra-frame features over the remaining frames that are not encoded yet. (4) The average of the inter-frame features averages the inter-frame features over the remaining frames. Note that we refer to the original frames for computing frame differences whenever the coded reference frames are not available. (5) The number of remaining bits in percentage terms evaluates $(R_{GOP}-R_{actual})/R_{GOP}$, where $R_{GOP}$ and $R_{actual}$ are the GOP-level rate constraint and the number of bits encoded up to the current frame, respectively. (6) The number of remaining frames in the GOP counts the frames that have not been encoded yet. (7) The temporal identification signals the level of the frame prediction hierarchy to which the current frame belongs. (8) The rate constraint of the GOP indicates the target bit rate for the current GOP. (9) The base QP specifies the default QP values for I-, B-, and b-frames, which are determined a priori irrespective of the input video.

The action of our RL agent indicates the QP difference (known as the delta QP) from the base QP. That is, the final QP value (i.e. $QP_{i}$ in Fig. 1) is the sum of the delta and the base QP’s. In our implementation, the delta QP ranges from -5 to 5 irrespective of the frame type (I-, B-, or b-frame). The fact that the agent learns to decide the delta QP helps reduce the search space. Note that different frame types have different base QPs (see Section 4.1).

The ultimate goal of RL is to maximize the expected cumulative reward in the long run. In order for the agent to behave in a desired way, it is crucial to define proper rewards.

In this work, we specify two immediate rewards, the distortion reward $r_{D_{i}}$ and the rate reward $r_{R_{i}}$. The former $r_{D_{i}}$ is defined as

where $\text{VMAF}_{i}$ and $\overline{\text{VMAF}}_{i}$ are the VMAF [7] scores of frame $i$ when encoded with the QP values chosen by our learned agent and by the rate control algorithm of x265, respectively. In other words, $r_{D_{i}}$ characterizes the gain in VMAF achieved by our policy. This gain is intended to be maximized.

The immediate rate reward, denoted by $r_{R_{i}}$, given to every video frame in a GOP is zero, except the last frame, the $r_{R_{i}}$ of which is given by

The sum $\sum_{i=1}^{N}r_{R_{i}}$ represents the negative absolute deviation of the GOP bit rate from the target $R_{GOP}$ in percentage terms.

With these immediate rewards, the distortion reward-to-go $Q_{D}(s,QP)$ and the rate reward-to-go $Q_{R}(s,QP)$ in Section 3.2 are evaluated as

where $\gamma=0.99$ is the discount factor. In particular, these reward-to-go functions are approximated by two separate networks, known as the distortion critic and the rate critic, respectively.

This section presents how we make use of the two critic networks to implement NFWPO. We begin with defining the feasible set $\mathcal{C}(s_{i})$ in coding a video frame $i$. In order to meet the GOP-level rate constraint, the feasible set $\mathcal{C}(s_{i})$ includes the QP values satisfying the requirement that the rate reward-to-go $Q_{R}$ is greater than or equal to a threshold $\epsilon$, given the current state $s_{i}$ (see Fig. 2 (a)):

According to Eqs. (9) and (11), this implies that $\mathcal{C}(s_{i})$ includes QP values that ensure the absolute rate deviation is upper bounded by the discounted $\epsilon$. In particular, we discretize these QP values by querying the rate critic $Q_{R}$ in discrete steps of 0.1 within the delta QP range (i.e. $QP_{i}=\{\text{base QP}-5,\text{base QP}-4.9,...,\text{base QP}+5$}).

Given the feasible set $\mathcal{C}(s_{i})$, we follow the process in Section 2 to generate the reference action $\tilde{a}_{s}$. Fig. 2 (b) illustrates the process, in which the actor output $\pi(s_{i})$ is assumed to be outside of the feasible set. As shown, $\pi(s_{i})$ will first be projected onto the feasible set to arrive at $\Pi_{\mathcal{C}(s_{i})}(\pi(s_{i}))$. We then evaluate $\bar{c}(s_{i})$ according to Eq. (3), where we use $Q_{D}(s,a)$ for $Q(s,a)$, and obtain $\tilde{a}_{s}$ based on Eq. (4). Finally, we update the actor network by Eq. (5).

Algorithm 1 details our proposed NFWPO-based RL algorithm. Similar to the DDPG algorithm, lines 5 to 12 roll out the learned policy with an exploration noise. The resulting transitions of states, actions, and rewards are stored in the replay buffer $R$ and are sampled to update the critics $Q_{D}$ and $Q_{R}$ in lines 15 and 17, respectively. Lines 19 to 22 correspond to NFWPO (Section 2).

Our actor and critic networks have similar network architectures to [5], but take different state signals as inputs.

We experiment with the proposed method on x265 using a hierarchical-B coding structure (GOP=16) as depicted in Fig. 3. To focus on the GOP-level bit allocation scheme, we follow the same GOP-level bit allocation as x265. To this end, we first encode every test sequence with fixed QP’s 22, 27, 32, and 37. This establishes four target bit rates $R_{s}$ at the sequence level. Next, we enable the 2-pass ABR rate control mode of x265 (–pass $2$ –bitrate $R_{s}$ –vbv-bufsize $2\times R_{s}$ –vbv-maxrate $2\times R_{s}$) in generating the anchor bitstreams, to meet these target bit rates. Finally, we observe the GOP-level bit rates of x265 and use them as our $R_{GOP}$. Since we use VMAF as the quality metric, to obtain the best perceptual quality, we turn off the –tune option in accordance with the x265 manual. It is to be noted that all the RL-based schemes (i.e. the single-critic method, the dual-critic method, and ours) are one-pass schemes. They do not make use of any information in the first encoding pass for bit allocation.

For training our model, the training dataset includes video sequences in UVG, MCL-JCV, and JCT-VC Class A. To expedite the training, all these sequences are resized to $512\times 320$. Likewise, the test videos, which are Class B and Class C sequences in JCT-VC dataset, are resized to the same spatial resolution. The single-critic and dual-critic methods use the same dataset for training.

We determine the hyper-parameters as follows. The base QP are selected as $QP_{l}-3$, $QP_{l}-2$, and $QP_{l}+2$ for I-, B-, and b-frame, respectively, where $QP_{l}$ are 22, 27, 32, and 37. The learning rate $\alpha$ for NFWPO in Eq. (4) is set to $0.1$, and the learning rate for the actor and critic networks is $0.001$. We use a discount factor $\gamma=0.99$ and the 3-step temporal difference learning to train our critics. The threshold $\epsilon$ of the feasible set is set to $-0.05$, allowing for a maximum rate deviation of $\pm 5\%$. For a fair comparison, the same rate requirement applies to the single- and dual-critic methods.

We assess the competing methods in terms of their R-D performance (in VMAF) and rate control accuracy.

Table 1 presents BD-rate results, with the 2-pass ABR of x265 serving as anchor. From Table 1, our scheme outperforms the single-critic [1] and the dual-critic [5] methods, both of which exhibit $3\%-4\%$ rate inflation as compared to x265 (operated in ABR mode with two-pass encoding). Ours achieves comparable performance to x265, yet with one-pass encoding. One point to note here is that all three RL-based models, including ours, perform poorly on BQMall. When this out-liner is excluded, our scheme performs slightly better than x265, showing a 0.41% average bit rate saving.

To shed light on why our scheme does not work well on BQMall, we overfit our NFWPO-based model to the test sequences. In other words, the model is trained on the test sequences. The overfitting results presented in the right most column of Table 1 indicate that BQMall, on which our previous model performs poorly, benefit the most from overfitting. Moreover, the overfit model performs comparably to the non-overfit model on the other test sequences. This then implies that our NFWPO-based RL algorithm and the model capacity are less of a problem. The root cause of the poor performance on BQMall may be attributed to the less representative training data.

Table 2 further presents the average bit-rate deviations from the GOP target $R_{GOP}$. The results are provided only for the highest and the lowest rate points due to the space limitation. In reporting the average rate deviation, any deviation from the $R_{GOP}$ within $\pm 5\%$ is regarded as $0\%$ to account for our $\pm 5\%$ tolerance margin. As seen from the table, our model shows fairly good rate control accuracy. The average rate deviations at high and low rates are around $2.1\%$ and $2.5\%$, respectively. In contrast, the rate deviations of the dual-critic method average $4.7\%$ and $4.9\%$ at high and low rates, respectively. The single-critic method has relatively inconsistent results: at high rates, the average rate deviation is as low as $1.2\%$, whereas at low rates, it is as high as $13.9\%$. This inconsistency in rate control accuracy between the low and high rates is attributed to the use of a fixed $\lambda$ for trading off the rate penalty against the distortion reward.

To understand the effectiveness of our rate control scheme, the top row of Fig. 4 visualizes the GOP-level bit rates for RaceHorses and BQMall while the bottom row presents the corresponding VMAF scores. The objective here is to achieve better VMAF results when encoding every GOP at the same bit rate as x265 (Section 4.1). As shown, applying the base QP alone, which corresponds to a fixed nested QP scheme, can hardly match the bit rates of x265. However, with our RL-based delta QP correction, the resulting bit rates match closely those of x265. Furthermore, our scheme shows similar or better VMAF scores than x265 in most GOPs, except for GOPs 12 to 24 in BQMall. This confirms the effectiveness of our RL training.

Fig. 5 presents the GOP-level bit rates, the frame-level QP assignment, and the corresponding VMAF scores for the competing methods. Again, the objective is to match the GOP-level bit rate of x265 while maximizing the VMAF scores (Section 4.1).

From the QP assignment, we see that the single-critic method tends to choose smaller QP values for I-frames, causing the resulting bit rate to exceed the target bit rate. The results presented correspond to the low-rate setting, where the single-critic method usually shows larger rate deviations (Section 4.2 & Table 2). We attribute the reason to the use of a fixed $\lambda$ hyper-parameter in formulating a single reward function $r_{D}+\lambda r_{R}$. The varying dynamics between $r_{D}$ and $r_{R}$ at different bit rates and on different sequences may render the rate reward/constraint less prominent in some cases. In contrast, the dual-critic method has fairly precise rate control accuracy. It, however, shows lower VMAF scores because of choosing higher QP values for I-frames. The general observation is that at low rates, the dual-critic method tends to disregard the distortion in order to meet the stricter rate constraint. Different from these two methods, our model chooses similar QP values to x265 for I-frames. An intriguing finding is that it tries to improve the VMAF score by choosing relatively smaller QP values for the B-frame and the first group of b-frames (frames 1-7 in Fig. 3) in a GOP. It then chooses larger QP values for the second group of b-frames (frames 9-15 in Fig. 3) to meet the rate constraint. Another observation is that our NFWPO-based model is relatively more robust to varied coding conditions, showing good rate control accuracy at both low and high rates.

Fig. 6 presents a subjective quality comparison. We see that our proposed method preserves more texture details and shows sharper image quality. More subjective quality comparisons are provided in the supplementary document.