Actor Critic using Kronecker-Factored Trust Region (ACKTR)

ACKTR gives actor-critic networks second-order updates without ever building the giant Fisher matrix. It splits each layer curvature into two small Kronecker factors that are cheap to invert.

ACKTR splits each layer Fisher block into two small Kronecker factors, turning one intractable inversion into two tiny ones

Why Does This Exist?

A plain policy gradient steps in weight space: change every weight by a fixed-size nudge along the gradient. The trouble is that weights are not the policy. A small nudge to a few weights can flip the action distribution from exploratory to near-deterministic, and the update has no idea it just did that. Training oscillates or collapses, and you pay for it in samples.

The natural gradient fixes the units. It measures distance between policies, not between weight vectors, using the Fisher information matrix $F$ . The update $\tilde{g} = F^{-1} g$ is the steepest direction in distribution space, so a step of a given size changes behavior by a predictable amount. Trust Region Policy Optimization uses this idea but never builds $F$ : it approximates $F^{-1} g$ with conjugate-gradient iterations, needs repeated curvature-vector passes, and applies the second-order treatment to the actor only while the critic trains with a first-order optimizer.

ACKTR (Wu et al., 2017) makes the exact second-order step practical for both networks. Its trick is Kronecker-factored approximate curvature (K-FAC): each layer Fisher block factors into two small covariance matrices, one from activations and one from output gradients. Inverting those costs almost nothing next to inverting $F$ itself.

This page covers the factorization, the trust-region step, and a hand-computed 2 by 2 example. For the trust-region idea underneath it, read the TRPO page first if you have not.

Think of It Like This

Tuning a grid of sprinklers by rows and columns

Picture a lawn with sprinklers in a 100 by 100 grid. Each sprinkler has its own valve, so there are 10,000 settings. You want to correct the watering pattern, and the effect of each valve depends on water pressure along its row and its column.

The naive approach measures how every valve interacts with every other valve: a 10,000 by 10,000 table of interactions. Nobody has time for that.

Instead you notice the structure. Row pressure varies independently of column pressure, to a good approximation. So you measure two small things: a 100 by 100 table of row effects and a 100 by 100 table of column effects. To correct any single valve, you combine its row correction with its column correction. Two small inversions replace one impossible one.

Where it stops: real rows and columns interact a little (a clogged main affects both), and K-FAC likewise pretends activations and output gradients are statistically independent when it splits the expectation. That approximation is the price of tractability, and in practice it is worth paying.

How It Actually Works

Take one dense layer mapping activations $a$ (length $d_{\text{in}}$ ) to pre-activations $s = W a$ (length $d_{\text{out}}$ ). Here $W$ is the weight matrix, $a$ is the layer input for one sample, and $g_s$ is the gradient of the loss with respect to $s$ (how much each pre-activation should change).

The gradient with respect to the weights is an outer product:

\nabla_W L = g_s a^\top

Flattened into a vector, that outer product becomes a Kronecker product, written $\otimes$ :

\text{vec}(\nabla_W L) = a \otimes g_s

The layer Fisher block is the expected outer product of these flattened gradients:

F_l = \mathbb{E}\left[ (a a^\top) \otimes (g_s g_s^\top) \right]

K-FAC splits this expectation into a product of expectations:

F_l \approx \mathbb{E}[a a^\top] \otimes \mathbb{E}[g_s g_s^\top] = A \otimes S

where $A$ is the $d_{\text{in}}$ by $d_{\text{in}}$ activation covariance and $S$ is the $d_{\text{out}}$ by $d_{\text{out}}$ output-gradient covariance. The inverse of a Kronecker product is the Kronecker product of the inverses, so with $G = \nabla_W L$ :

\tilde{G} = S^{-1} G A^{-1}

That is the whole trick. Instead of inverting a $(d_{\text{in}} d_{\text{out}})$ by $(d_{\text{in}} d_{\text{out}})$ matrix, you invert one small matrix per side and multiply twice. For a 100 by 100 layer that is two 100 by 100 inversions instead of one 10,000 by 10,000 inversion.

The trust-region step

The natural gradient sets the direction; a trust region sets the size. With $\tilde{g}$ the flattened natural gradient, $g$ the plain gradient, and $\delta$ the allowed KL change per update (usually about 0.01):

\beta = \min\left(1, \sqrt{\frac{2\delta}{\tilde{g}^\top g}}\right), \qquad \Delta \theta = \beta \, \tilde{g}

If the proposed step would move the policy too far, $\beta$ shrinks it. In practice the covariances $A$ and $S$ are tracked as running averages over minibatches, a small damping term $\gamma I$ is added before inverting for stability, and the inverses are recomputed every few dozen steps rather than every step.

Worked example

A 2 by 2 layer. Suppose the tracked covariances (damping already folded in) are diagonal:

A = \begin{bmatrix} 2.0 & 0.0 \\ 0.0 & 1.0 \end{bmatrix}, \qquad S = \begin{bmatrix} 4.0 & 0.0 \\ 0.0 & 2.0 \end{bmatrix}, \qquad G = \begin{bmatrix} 1.0 & 1.0 \\ 1.0 & 1.0 \end{bmatrix}

Invert each side: $A^{-1} = \text{diag}(0.5, 1.0)$ and $S^{-1} = \text{diag}(0.25, 0.5)$ . First multiply $S^{-1} G$ :

S^{-1} G = \begin{bmatrix} 0.25 & 0.25 \\ 0.50 & 0.50 \end{bmatrix}

Then multiply by $A^{-1}$ on the right, which scales column one by 0.5 and leaves column two alone:

\tilde{G} = \begin{bmatrix} 0.125 & 0.25 \\ 0.25 & 0.50 \end{bmatrix}

Check against the full Kronecker inverse: $A \otimes S = \text{diag}(8.0, 4.0, 4.0, 2.0)$ , whose inverse is $\text{diag}(0.125, 0.25, 0.25, 0.50)$ . Applied to the flattened gradient $[1, 1, 1, 1]$ , that gives $[0.125, 0.25, 0.25, 0.50]$ , which reshapes to exactly $\tilde{G}$ . The sandwich agrees with the giant inverse.

Now cap the step with $\delta = 0.01$ . The inner product $\tilde{g}^\top g = 0.125 + 0.25 + 0.25 + 0.50 = 1.125$ , so:

\beta = \min\left(1, \sqrt{0.02 / 1.125}\right) \approx 0.13333

The final update is $0.13333 \times \tilde{G}$ :

\Delta W = \begin{bmatrix} 0.0167 & 0.0333 \\ 0.0333 & 0.0667 \end{bmatrix}

Notice how the curvature reshaped the step: the plain gradient treated all four weights equally, but the natural gradient pushes hardest on the weight whose activation and gradient variances are both smallest.

Code

This snippet computes the K-FAC sandwich and the trust-region scale for the worked example above, in pure standard-library Python. Its point is the check: the fast sandwich form must match the full Kronecker inverse.

1from typing import List2
3
4def mat_mul(a: List[List[float]], b: List[List[float]]) -> List[List[float]]:5    """Plain matrix product for small dense matrices."""6    n, m, p = len(a), len(b[0]), len(b)7    return [[sum(a[i][k] * b[k][j] for k in range(p)) for j in range(m)] for i in range(n)]8
9
10def diag_inv(d: List[float]) -> List[List[float]]:11    """Inverse of a diagonal matrix given as a list of diagonal entries."""12    return [[1.0 / d[i] if i == j else 0.0 for j in range(len(d))] for i in range(len(d))]13
14
15def kfac_step(16    grad: List[List[float]],17    act_diag: List[float],18    grad_diag: List[float],19    delta: float = 0.01,20) -> tuple:21    """Return (natural gradient, trust-region step, scale) for one dense layer."""22    nat = mat_mul(mat_mul(diag_inv(grad_diag), grad), diag_inv(act_diag))23    curve = sum(g * n for grow, nrow in zip(grad, nat) for g, n in zip(grow, nrow))24    beta = min(1.0, (2.0 * delta / curve) ** 0.5)25    step = [[beta * v for v in row] for row in nat]26    return nat, step, beta27
28
29nat, step, beta = kfac_step([[1.0, 1.0], [1.0, 1.0]], [2.0, 1.0], [4.0, 2.0])30print([[round(v, 4) for v in row] for row in nat])31# -> [[0.125, 0.25], [0.25, 0.5]]32print(f"{beta:.5f}")33# -> 0.1333334print([[round(v, 4) for v in row] for row in step])35# -> [[0.0167, 0.0333], [0.0333, 0.0667]]

All three prints match the hand computation, including the trust-region scale.

Watch Out For

Inverting every step, or damping it wrong

Inverting $A$ and $S$ on every minibatch is slow, and the inverses go stale at different rates than the gradients. The standard fix is to refresh the running covariances every batch but recompute the inverses only every 20 to 100 steps. If training suddenly diverges after running fine, a stale inverse is a prime suspect: refresh them and see if the spike disappears.

Damping cuts the other way. Too little ( $\gamma$ near zero) and a near-singular covariance produces an enormous step out of nowhere. Too much and $(A + \gamma I)^{-1}$ is nearly a scaled identity, so you paid for second-order machinery and got plain gradient descent back. When the natural step stops helping, check $\gamma$ before blaming the idea.

Assuming the factorization fits every layer

K-FAC in its basic form assumes a dense layer with no weight sharing. Convolutional layers share weights across spatial positions, so the clean $A \otimes S$ split needs patch-based statistics instead of plain activation covariances. Normalization layers and recurrent weights break the independence assumption further. If you apply the dense-layer formula verbatim to a conv net and updates explode, the layer type is the bug, not your arithmetic.

The Quick Version

The natural gradient $\tilde{g} = F^{-1} g$ steps in policy space instead of weight space, so fixed-size steps change behavior predictably.
K-FAC approximates each layer Fisher block as $A \otimes S$ , two small covariances, and the natural gradient becomes the sandwich $S^{-1} G A^{-1}$ .
A 100 by 100 layer needs two 100 by 100 inversions instead of one 10,000 by 10,000 inversion.
ACKTR applies this to both actor and critic, unlike TRPO, and needs no conjugate-gradient loop.
The trust-region scale $\beta = \min(1, \sqrt{2\delta / \tilde{g}^\top g})$ caps each step at a KL budget of about 0.01.
The worked 2 by 2 case gives $\tilde{G} = [[0.125, 0.25], [0.25, 0.5]]$ with $\beta \approx 0.13333$ .