av1enc: fix the bitrate overshoot of the bitrate mode rate control

Bitrate mode gave key frames keyint-dependent extra bits (about 6.8 frame
budgets at the 60 frame default) but left the inter frames at the full
frame budget, so every group of frames overshot by (boost - 1) / keyint,
9.6 % at keyint 60 (up to +187 % at keyint 5 on some content). The
accumulated deviation was tracked against the boosted budget, so the
feedback never saw it.

Now a group of keyint frames adds up to keyint * frame budget: inter
frames get (keyint - boost) / (keyint - 1) of it. Further changes to the
frame level control:

- The quantizer follows the size error with a measured slope (quantizer
  steps per halving of the size, taken from consecutive inter frames)
  instead of a fixed 12 steps per doubling. On noisy content the curve is
  several times steeper than that and the old gain oscillated between
  1 kB and 42 kB frames.
- Key frames have their own quantizer offset, adapted from key frame
  sizes only, instead of feeding the key frame error into the inter
  quantizer. A key frame after the first that misses its budget by more
  than 19 % is coded again (at most twice more), with a corrected
  quantizer found in the log domain of the quantizer step.
- A deviation from the budget is paid back over 16 frames, bounded to
  0.35x..2x of the budget.

Constant quantizer mode is untouched (byte identical output). The
encoded stream in bitrate mode changes.

Adds a test that the rate stays within 5 % of the target over 180
frames for three 8 bit and one 10 bit setting.

Part of #2391

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/av1enc/README.md

@@ -48,8 +48,10 @@ reconstruction, so prediction never drifts.
   reconstruction loop; full coefficient coding (all_zero, tx type, eob_pt, eob_extra,
   coeff_base_eob, coeff_base, coeff_br, dc_sign, Golomb) with the exact contexts.
 - **Rate control**: `ConstantQ(base_q_idx)`; `Bitrate { kbps }` adapts `base_q_idx` per frame
-  towards the average target (key frames get a lower quantizer and a larger share; the first
-  key frame is searched in a few passes).
+  towards the average target (key frames get a lower quantizer and a larger share of the budget of
+  their group of `keyint` frames, the inter frames the rest; the first key frame is searched in a
+  few passes, a later key frame that misses its budget by more than 19 % is coded again; the
+  deviation from the budget is paid back over 16 frames).
 
 ## Limits
 

crates/av1enc/src/lib.rs

@@ -24,6 +24,7 @@ use bitw::{BitWriter, obu};
 use decide::{Decider, FrameGeom, RefFrame};
 use loopfilter::{LfInfo, loop_filter};
 use pred::Plane;
+use tables::AC_QLOOKUP;
 use writer::TileWriter;
 
 /// Rate control.
@@ -136,10 +137,109 @@ struct RateState {
     frame_bits: f64,
     /// Quantizer index for inter frames (continuous).
     q: f64,
+    /// Key frame quantizer index minus the inter frame one.
+    key_off: f64,
+    /// Quantizer index steps per halving of the frame size, measured from consecutive inter
+    /// frames (the size/quantizer curve is much steeper than the default on noisy content).
+    slope: f64,
+    /// Quantizer index and size (in bits) of the previous inter frame.
+    prev_inter: Option<(f64, f64)>,
     /// Accumulated (spent - budget) bits.
     debt: f64,
 }
 
+/// Quantizer index steps per halving of the frame size before any measurement.
+const RC_SLOPE: f64 = 28.0;
+/// Fraction of the quantizer step the size error asks for that is applied per frame.
+const RC_GAIN: f64 = 0.8;
+/// Frames over which an accumulated deviation from the budget is paid back.
+const RC_PAYBACK: f64 = 16.0;
+
+impl RateState {
+    /// The nominal bit budget and the starting quantizer index of a key or inter frame. A key
+    /// frame gets a larger share, and the inter frames of the group get the rest, so that a whole
+    /// group of `keyint` frames adds up to `keyint * frame_bits`.
+    fn frame_budget(&self, keyint: u32, key: bool) -> (f64, f64) {
+        let keyint = keyint.max(1) as f64;
+        let boost = (1.0 + (keyint.clamp(1.0, 60.0) - 1.0).sqrt() * 0.75).min(0.5 * keyint + 0.5);
+        if key {
+            (self.frame_bits * boost, (self.q + self.key_off).clamp(1.0, 255.0))
+        } else if keyint > 1.0 {
+            (self.frame_bits * (keyint - boost) / (keyint - 1.0), self.q)
+        } else {
+            (self.frame_bits, self.q)
+        }
+    }
+
+    /// Adapts the quantizers to the size of the frame just coded at `q_used` against `target`.
+    fn update(&mut self, key: bool, first: bool, q_used: f64, bits: f64, target: f64) {
+        let err = (bits / target).log2().clamp(-1.5, 1.5);
+        if first && key {
+            // four passes ran: the key frame quantizer is known, the inter one follows from the offset
+            self.q = (q_used - self.key_off).clamp(1.0, 255.0);
+        } else if key {
+            self.key_off = q_used - self.q + RC_GAIN * self.slope * err;
+        } else {
+            if let Some((pq, pbits)) = self.prev_inter {
+                let (dq, dr) = (q_used - pq, (pbits / bits).log2());
+                if dq.abs() >= 3.0 && dr.abs() >= 0.15 && dq * dr > 0.0 {
+                    self.slope = (0.6 * self.slope + 0.4 * (dq / dr)).clamp(3.0, 45.0);
+                }
+            }
+            self.prev_inter = Some((q_used, bits));
+            self.q = (q_used + RC_GAIN * self.slope * err).clamp(1.0, 255.0);
+        }
+    }
+}
+
+/// A key frame after the first is coded again when its size is more than this many octaves off
+/// the target (19 %).
+const RC_KEY_TOLERANCE: f64 = 0.25;
+/// Octaves of frame size per octave of quantizer step assumed for the first correction of a frame.
+const RC_PROBE_SLOPE: f64 = 1.0;
+/// Largest change of the quantizer step (octaves) between two passes on one frame.
+const RC_PROBE_MAX_STEP: f64 = 2.0;
+
+/// The quantizer index whose step size is closest to `2^x`.
+fn q_for_step(bit_depth: u8, x: f64) -> u8 {
+    let table = &AC_QLOOKUP[((bit_depth - 8) / 2) as usize];
+    let mut best = 1;
+    for q in 1..=255usize {
+        if ((table[q] as f64).log2() - x).abs() < ((table[best] as f64).log2() - x).abs() {
+            best = q;
+        }
+    }
+    best as u8
+}
+
+/// The next quantizer to try for a frame, from the `(quantizer, log2(size / target))`
+/// pairs coded so far. The size follows the quantizer step (not the index, which is linear in the
+/// step only below about 60) in a near log-log line: false position between the closest pair of
+/// passes that enclose the target, otherwise a step from the latest one with the slope of the
+/// last two. `None` when no untried quantizer lies between the enclosing passes.
+fn next_probe_q(bit_depth: u8, tried: &[(u8, f64)]) -> Option<u8> {
+    let table = &AC_QLOOKUP[((bit_depth - 8) / 2) as usize];
+    let x = |q: u8| (table[q as usize] as f64).log2();
+    let over = tried.iter().filter(|t| t.1 > 0.0).max_by_key(|t| t.0);
+    let under = tried.iter().filter(|t| t.1 <= 0.0).min_by_key(|t| t.0);
+    if let (Some(&(q_lo, e_lo)), Some(&(q_hi, e_hi))) = (over, under) {
+        if q_hi - q_lo < 2 {
+            return None;
+        }
+        let next = x(q_lo) + (x(q_hi) - x(q_lo)) * e_lo / (e_lo - e_hi);
+        return Some(q_for_step(bit_depth, next).clamp(q_lo + 1, q_hi - 1));
+    }
+    let &(q, err) = tried.last()?;
+    let mut slope = RC_PROBE_SLOPE;
+    if let [.., (q0, e0), (q1, e1)] = tried
+        && (x(*q1) - x(*q0)).abs() > 0.05
+        && (e0 - e1) / (x(*q1) - x(*q0)) > 0.0
+    {
+        slope = ((e0 - e1) / (x(*q1) - x(*q0))).clamp(0.25, 3.0);
+    }
+    Some(q_for_step(bit_depth, x(q) + (err / slope).clamp(-RC_PROBE_MAX_STEP, RC_PROBE_MAX_STEP)))
+}
+
 /// AV1 encoder. Frames come out in input order (no hidden/reordered frames).
 pub struct Encoder {
     cfg: EncoderConfig,
@@ -185,7 +285,7 @@ impl Encoder {
                 // initial guess from the bits per sample
                 let bpp = frame_bits / (cfg.width as f64 * cfg.height as f64);
                 let q = (110.0 - 40.0 * (bpp / 0.1).log2()).clamp(20.0, 230.0);
-                Some(RateState { frame_bits, q, debt: 0.0 })
+                Some(RateState { frame_bits, q, key_off: -24.0, slope: RC_SLOPE, prev_inter: None, debt: 0.0 })
             }
         };
         let mut enc = Encoder { cfg, geom, level, seq_obu: Vec::new(), frame_num: 0, refr: None, rc };
@@ -469,38 +569,52 @@ impl Encoder {
 
     fn encode_rate_controlled(&mut self, src: &[Plane; 3], key: bool, rc: &mut RateState) -> (Vec<u8>, RefFrame) {
         let first = self.frame_num == 0;
-        let (frame_bits, debt, q_inter) = (rc.frame_bits, rc.debt, rc.q);
-        // key frames get a lower quantizer and a larger share of the budget
-        let key_boost = 1.0 + (self.cfg.keyint.clamp(1, 60) as f64 - 1.0).sqrt() * 0.75;
-        let budget = if key { frame_bits * key_boost } else { frame_bits };
-        // pay back a fraction of the accumulated deviation
-        let target = (budget - debt * 0.25).max(budget * 0.25);
-        let mut q = if key { q_inter - 24.0 } else { q_inter };
-        let mut best: Option<(Vec<u8>, RefFrame, f64)> = None;
-        let passes = if first { 4 } else { 1 };
-        for pass in 0..passes {
+        let (budget, q_start) = rc.frame_budget(self.cfg.keyint, key);
+        // pay back the accumulated deviation over the next RC_PAYBACK frames
+        let target = (budget - rc.debt / RC_PAYBACK).clamp(budget * 0.35, budget * 2.0);
+        let mut q = q_start;
+        let mut best: Option<(Vec<u8>, RefFrame, f64, f64)> = None;
+        // The first frame is coded four times to find its quantizer. A later key frame is coded
+        // again (up to twice) if its size is far off, since a key frame is a large part of the
+        // budget of its group.
+        let (max_passes, tolerance) = if first {
+            (4, 0.0)
+        } else if key {
+            (3, RC_KEY_TOLERANCE)
+        } else {
+            (1, 0.0)
+        };
+        let mut tried: Vec<(u8, f64)> = Vec::new();
+        for pass in 0..max_passes {
             let qi = q.round().clamp(1.0, 255.0) as u8;
+            if !first && tried.iter().any(|t| t.0 == qi) {
+                break;
+            }
             let (obu, rf) = self.encode_frame(src, key, qi);
             let bits = obu.len() as f64 * 8.0;
             let err = (bits / target).log2();
-            let better = best.as_ref().is_none_or(|b| err.abs() < (b.2 / target).log2().abs());
-            if better {
-                best = Some((obu, rf, bits));
+            tried.push((qi, err));
+            if best.as_ref().is_none_or(|b| err.abs() < (b.2 / target).log2().abs()) {
+                best = Some((obu, rf, bits, qi as f64));
             }
-            if pass + 1 < passes {
-                q = (q + 30.0 * err).clamp(1.0, 255.0);
+            if err.abs() <= tolerance {
+                break;
+            }
+            if first {
+                if pass + 1 < max_passes {
+                    q = (q + 30.0 * err).clamp(1.0, 255.0);
+                }
+            } else {
+                match next_probe_q(self.cfg.bit_depth, &tried) {
+                    Some(next) => q = next as f64,
+                    None => break,
+                }
             }
         }
         // At least one pass ran, so `best` is set.
-        let Some((obu, rf, bits)) = best else { return self.encode_frame(src, key, q.round().clamp(1.0, 255.0) as u8) };
+        let Some((obu, rf, bits, q_used)) = best else { return self.encode_frame(src, key, q.round().clamp(1.0, 255.0) as u8) };
         rc.debt += bits - budget;
-        // adapt the inter quantizer: about 30 index steps per doubling of the rate
-        let err = (bits / target).log2().clamp(-1.0, 1.0);
-        if first && key {
-            rc.q = (q + 24.0).clamp(1.0, 255.0);
-        } else {
-            rc.q = (rc.q + 12.0 * err).clamp(1.0, 255.0);
-        }
+        rc.update(key, first, q_used, bits, target);
         (obu, rf)
     }
 

crates/av1enc/tests/it/roundtrip.rs

@@ -80,6 +80,23 @@ fn bitrate_mode_tracks_target() {
     assert!(kbps > 400.0 / 2.0 && kbps < 400.0 * 2.0, "{kbps} kbit/s");
 }
 
+/// Over several key frame intervals the average rate stays within 5 % of the target, for a
+/// range of targets (the key frame budget is taken out of the inter frames of its group).
+#[test]
+fn bitrate_mode_is_accurate_over_many_groups() {
+    let n = 180;
+    for (kbps, keyint, bd) in [(40, 30, 8), (100, 30, 8), (200, 60, 8), (150, 20, 10)] {
+        let pics: Vec<Pic> = (0..n).map(|t| synth(192, 128, bd, t)).collect();
+        let mut c = cfg(192, 128, bd as u8, 0, keyint);
+        c.rate = RateControl::Bitrate { kbps };
+        let mut enc = effectcraft_av1enc::Encoder::new(c).unwrap();
+        let bits: usize = pics.iter().map(|p| enc.encode(&p.frame()).data.len() * 8).sum();
+        let got = bits as f64 / (n as f64 / 30.0) / 1000.0;
+        eprintln!("{bd} bit, keyint {keyint}: {got:.1} kbit/s for {kbps}");
+        assert!((got / kbps as f64 - 1.0).abs() < 0.05, "{got:.1} kbit/s for a {kbps} kbit/s target ({bd} bit, keyint {keyint})");
+    }
+}
+
 #[test]
 fn headers() {
     let enc = effectcraft_av1enc::Encoder::new(cfg(1920, 1080, 10, 100, 60)).unwrap();