PolicyAttention: Softmax Attention Implements Policy Mirror Descent for Closed-Loop Control

Yuhe Sui; Yingzhi Tang; Shufang Chen

cs.LGSep 24, 2026

PolicyAttention: Softmax Attention Implements Policy Mirror Descent for Closed-Loop Control

Authors: Yuhe Sui, Yingzhi Tang, Shufang Chen

Organizations: Quantitative Research Society, Singapore · Nanyang Technological University, Singapore · The Chinese University of Hong Kong, Shenzhen, China · The University of Hong Kong, Hong Kong SAR, China

Abstract

Can causal softmax attention implement policy mirror descent as a repeated controller rather than a one-step algebraic identity? Negative-entropy policy mirror descent (PMD) has the statewise update $\operatorname{PMD}_η(π,Q)=\operatorname{softmax}(\logπ+ηQ)$ . Building on the known Q-TD-PMD recursion, we construct one fixed causal-softmax actor--environment--one-step-critic protocol with explicit actor, routing, sampling, and normalization residuals, and propagate them to the policy actually returned. The construction states the finite-logit/full-support domain, the external tokenization and sampling boundary, and the mean-zero LayerNorm carrier conditions required by the normalized compilation. Separately trained pre-LN Transformers recover the target computation empirically. A frozen one-step audit model is closest to PMD among the tested fixed rules; in a preregistered five-run $S=4$ repeated-control test, the learned actor with an exact one-step critic reaches median returned-policy loss $1.052\times$ the Exact PMD oracle and retains the criterion across four no-retraining shifts. The same checkpoints with their learned critic give descriptive median $1.050\times$ the oracle (no registered margin). At $S=8$ , replacing the exact critic by the learned critic raises median $T=20$ loss to $0.0225$ yet leaves the Liang--Lai and Algorithm Distillation adaptations $20.2$ -- $24.2\times$ higher-loss; this is a one-sided sampled-critic bound because PolicyAttention consumes 144 generative transitions per round versus 20 on-policy transitions for the adaptations. The strict 20-transition comparison remains open. At $S=8,16$ , the exact-critic common-harness comparison remains $17.7$ -- $28.2\times$ lower-loss than those adaptations, with the information asymmetry stated locally.

Figures & tables

Panel of Figure 1 — Figure 1: One mechanism-to-control chain. The causal loop is shown above; the table below separates the constructive theorem, the learned-controller evidence, and the literature-facing benchmark. Exact PMD is the non-learned oracle, while the learned PMD-target model serves only as calibration.

Work	Softmax	Eval.	Control	Learned	Return-error
Wang et al. 2025	emp.	✓	–	✓	–
Xie et al. 2026	✓	✓	–	✓	–
Liang–Lai 2026	linear	✓	✓	✓	–
Algorithm Distill.	✓	implicit	✓	✓	–
PolicyAttention	✓	✓	✓	✓	✓

Table 1: Source-verified capability map. “Softmax” denotes standard softmax attention in the relevant construction/model; PolicyAttention’s final column denotes the returned-policy residual interface introduced here.

Figure 2: The learned actor exhibits the predicted PMD signature before closed-loop testing. (a) On prior held-out one-step inputs, PMD has mean row-

L_{1}

error

0.0357

, versus

0.1116

for the nearest fixed alternative. (b) Refitted effective step size changes monotonically with supplied

\eta

, with shrinkage toward the training-range center. These are local mechanism measurements, not returned-policy performance.

Question	Coordinate / comparison	Verified result and interpretation
What actor rule was learned?	same-context row- $L_{1}$ to fixed update hypotheses	PMD $0.0357$ , nearest alternative $0.1116$ : strong local PMD correspondence on held-out one-step inputs.
Does local correspondence survive feedback?	learned actor + exact one-step critic / Exact PMD oracle	median ratio $1.052$ ; all five trained runs satisfy the preregistered $1.5$ criterion.
Does the fully learned loop remain near the oracle?	learned actor + learned critic / Exact PMD oracle	descriptive median ratio $1.050$ at $S=4$ ; no reliability margin was registered for this comparison.
Does it survive environment shift?	same learned actors + exact critic, four no-retraining families	all four families satisfy the declared criterion separately; the first three use 64 MDPs each; the structured ring is one MDP $\times$ 64 initial policies.
Does the phenomenon persist at larger size?	retrained $S=8$ and $S=16$ models	PolicyAttention remains within 1.31 $\times$ and 1.64 $\times$ the Exact PMD oracle at $T=20$ ; this is retraining, not zero-shot size generalization.
How does it compare with neighboring learned ICRL?	qualified common-harness adaptations	exact-critic PolicyAttention is $17.7$ – $28.2\times$ lower-loss at $T=20$ ; at $S=8$ the learned-critic check remains $20.2$ – $24.2\times$ lower but uses $7.2\times$ the adaptations sampling budget.

Table 2: Claim-to-estimand map. Each row answers a different scientific question; evidence on one row is not substituted for another.

Figure 3: Local PMD correspondence remains useful under repeated deployment. (a) With an exact one-step critic, all five confirmatory

S=4

runs stay within the preregistered

1.5\times

margin to the Exact PMD oracle. (b) The same trained actors satisfy the criterion across four no-retraining shifts; whiskers show observed five-run ranges.

Figure 4: PolicyAttention has substantially lower

T=20

returned-policy loss than the two qualified external adaptations. Markers show medians at

S=8,16

and small open points show trained-run medians. PolicyAttention uses an exact one-step model-based critic here, while the external adaptations receive 20 sampled transitions per round; this is a complete-controller comparison, not an information-matched actor comparison.

	$S=8$		$S=16$
Method	$T20$	$T40$	$T20$	$T40$
PolicyAttention	.0193	.0108	.0273	.0224
PMD-target ref. (learned)	.0189	.0106	.0232	.0144
PolicyAttention, learned critic	.0225	.0121	–	–
Liang–Lai adapt.	.4535	.2812	.4817	.3307
AD adapt.	.5439	.3950	.5723	.3730

Table 3: Median normalized returned-policy loss in the common harness (lower is better). Every learned method has three training identities. External entries are our adaptations, not the papers’ original benchmarks.

Method	actor/training signal	online access per round	parameters / training data	recorded training cost
PolicyAttention	PMD variational actor	exact one-step backup; 0 samples	272k; 24 MDPs, 2,048 examples, 51.2k steps	0.19–0.24 GPU-h
PMD-target ref. (learned)	direct Exact-PMD labels	exact one-step backup; 0 samples	same architecture/data/budget as PolicyAttention	0.19–0.24 GPU-h
Liang–Lai adaptation	linear-attention actor–critic teacher mimicking	20 on-policy transitions	10.9k; 10,000 MDPs, $10^{7}$ windows	about 17 CPU-min
AD adaptation	action NLL on source learning histories	20 sampled transitions	1.14M; 2,000 histories $\times$ 1,000 source steps	0.63–0.75 GPU-h

Table 4: Method interface and training context at

S=8

. “Exact backup” is one Bellman policy-evaluation step with model access, not full evaluation to

Q^{\pi}

. The learned-critic PolicyAttention check uses 144 generative transitions per round and is reported separately in text/Table 3 .

Figure 5: PolicyAttention occupies the low-loss, low-deviation region. (a) Same-context row-

L_{1}

deviation from Exact PMD versus

T=20

returned-policy loss at

S=8

and

S=16

; the dashed curve is the non-dominated envelope over the two external adaptations. (b) Common-harness

T=20

loss normalized by the learned PMD-target reference.

Explore similar work

Fast Regularized Policy Mirror Descent with One-Step TD Updates
Sep 30, 2026Qipei Chen, Wenye Li, Yule Sun +1Markov Decision Processes Policy Optimization
School of Mathematical Sciences, Fudan University · School of Data Science, Fudan University
Sharp Convergence and Sample Complexity of Policy Mirror Descent for Average-Reward MDPs
Oct 2, 2026Enes Arda, Atilla EryilmazMarkov Decision Processes Average-Reward RL
The Ohio State University
An analysis of Mirror-Descent Soft Actor-Critic
Sep 28, 2026Denis Zorba, Michal ValkoReinforcement Learning Online Mirror Descent
School of Mathematics, University of Edinburgh, United Kingdom · INRIA, Paris, France

Sep 30, 2026J/K move · Enter open · S save

Fast Regularized Policy Mirror Descent with One-Step TD Updates

Qipei Chen, Wenye Li, Yule Sun, Ke Wei

School of Mathematical Sciences, Fudan University · School of Data Science, Fudan University

Markov Decision Processes Policy Optimization Temporal-Difference Learning Online Mirror Descent Policy Learning Sample-Efficient RL

Open paper AskPDF arXiv

Policy mirror descent (PMD) enjoys fast convergence in regularized Markov decision processes (MDPs), but existing guarantees often rely on exact or increasingly accurate policy evaluation. We analyze PMD coupled with a persistent critic advanced by one temporal-difference (TD) update. For finite discounted MDPs, we establish global linear convergence in value for exact coordinate-wise Bellman updates, with any positive constant actor stepsize and arbitrary finite critic initialization. The proof combines a resolvent-based auxiliary distribution with a decaying Bellman-violation correction and a potential weighted by inverse coordinate weights. We then study stochastic TD-PMD with general strongly convex mirror maps under a single off-policy Markov trajectory. With suitably chosen constant stepsizes and a finite-batch TD update, the method achieves an expected value gap of

ε

after

\widetilde{O}(1/((1-γ)^5 \widetildeσ_b ε))

transitions. The stochastic analysis relies on the trajectory-wise Lipschitz continuity of the regularizer, derived from uniform bounds on vertex Bregman divergences, together with a visitation-weighted resolvent estimate for signed critic-error propagation that yields an inverse-linear dependence on behavior coverage

\widetildeσ_b

. In contrast to many prior guarantees for regularized policy optimization, our sample-complexity guarantee holds without trajectory resets, generative-model access, or nested policy-evaluation loops. Numerical results are consistent with the theoretical convergence analysis.

Abstract

Explore similar work

Fast Regularized Policy Mirror Descent with One-Step TD Updates

Sharp Convergence and Sample Complexity of Policy Mirror Descent for Average-Reward MDPs

An analysis of Mirror-Descent Soft Actor-Critic