Arithmetic Actor Heads and Training Stabilization for Out-of-Distribution Reinforcement Learning

Yifan Zhang; Liang Zheng

cs.LGOct 4, 2026

Arithmetic Actor Heads and Training Stabilization for Out-of-Distribution Reinforcement Learning

Authors: Yifan Zhang, Liang Zheng

Organizations: Central South University

Abstract

Reinforcement learning (RL) policies can deteriorate under out-of-distribution (OOD) magnitude shifts. Starting from soft actor-critic (SAC) and its Bayesian Amnesic Piecewise-Robust (BAPR) predecessor, we study the causal-symbolic BAPR (CS-BAPR) family. The practical method combines six training-stabilization settings with alternative actor heads: a Neural Addition Unit (NAU) with a Neural Multiplication Unit (NMU)-inspired quadratic correction, a Kolmogorov-Arnold Network (KAN), or a multilayer perceptron (MLP) with rectified linear unit (ReLU) or hyperbolic-tangent activations.

Figures & tables

Method	$N_{\mathrm{seed}}$ seeds	ID total return (K)
bapr-pristine (ReLU $+$ no fixes)	10	$-850\pm 37$
bapr (historical legacy)	10	$-793\pm 34$
csbapr-no-nau (ReLU-MLP $+$ 6 fixes)	10	$-780\pm 55$
csbapr (NAU $+$ 6 fixes)	10	$-772\pm 36$
csbapr-tanh (smooth-MLP $+$ 6 fixes)	10	$-817\pm 54$
csbapr-kan (KAN $+$ 6 fixes)	10	$-759\pm 35$

Table 1: ID total return across the 12 lines at

m=1

. Higher is better. Mean

\pm

standard deviation,

N_{\mathrm{seed}}{=}10

seeds per row.

Method	$m{=}1$	$m{=}2$	$m{=}5$	$m{=}10$	$m{=}20$	$m{=}50$
bapr-pristine (ReLU, no fixes)	$-861$	$-862$	$-939$	$-1134$	$-1267$	$-1416$
bapr (historical legacy)	$-791$	$-799$	$-878$	$-1025$	$-1196$	$-1428$
csbapr-no-nau (ReLU-MLP $+$ 6 fixes)	$-770$	$-776$	$-827$	$\mathbf{-919}$	$\mathbf{-1076}$	$-1314$
csbapr (NAU $+$ 6 fixes)	$-772$	$-778$	$\mathbf{-807}$	$-950$	$-1098$	$-1293$
csbapr-tanh (smooth-MLP $+$ 6 fixes)	$-786$	$-824$	$-897$	$-986$	$-1087$	$\mathbf{-1239}$
csbapr-kan (KAN $+$ 6 fixes)	$\mathbf{-756}$	$\mathbf{-762}$	$-820$	$-971$	$-1121$	$-1335$

Table 2: Static OOD demand sweep, total episode return (K, higher is better). Median episode return over

N_{\mathrm{seed}}{=}10

seeds per row.

m=1

is ID. bapr-pristine disables all six stabilization settings; bapr is a historical legacy configuration. Bold : best in column.

Method	burst $5{\times}$	burst $10{\times}$	burst $20{\times}$	burst $50{\times}$
bapr-pristine	$-909$	$-1021$	$-1178$	$-1354$
bapr (legacy)	$-836$	$-963$	$-1104$	$-1316$
csbapr-no-nau	$-792$	$\mathbf{-858}$	$\mathbf{-982}$	$-1205$
csbapr (NAU)	$-792$	$-881$	$-1000$	$-1196$
csbapr-tanh	$-847$	$-932$	$-1017$	$\mathbf{-1152}$
csbapr-kan	$\mathbf{-790}$	$-909$	$-1054$	$-1224$

Table 3: Abrupt-burst protocol:

m=1

for

t\in[0,3600)

s, then abrupt switch to the specified multiplier. Total episode return (K), median over

N_{\mathrm{seed}}{=}10

seeds.

Method	commuter_day	square_wave_20x	escalating_10x
bapr-pristine	$-998$	$-1111$	$-956$
bapr (legacy)	$-954$	$-1045$	$-920$
csbapr-no-nau	$\mathbf{-859}$	$\mathbf{-921}$	$\mathbf{-845}$
csbapr (NAU)	$-865$	$-927$	$-853$
csbapr-tanh	$-929$	$-990$	$-908$
csbapr-kan	$-914$	$-972$	$-912$

Table 4: Oscillating OOD protocol: within-episode piecewise-constant demand schedule. Total episode return (K), median over

N_{\mathrm{seed}}{=}10

seeds. Schedules: commuter_day has 5 switches with peak

50{\times}

; square_wave_20x has 9 switches with peak

20{\times}

; escalating_10x has 4 switches with peak

50{\times}

Method	$N_{\mathrm{seed}}$	$<{-}1{,}200$ K	$<{-}1{,}500$ K	$<{-}1{,}800$ K	$<{-}2{,}000$ K	Worst seed (K)
bapr-pristine (ReLU, no fixes)	10	$70\%$	$0\%$	$0\%$	$0\%$	$-1{,}422$
bapr (historical legacy)	10	$50\%$	$0\%$	$0\%$	$0\%$	$-1{,}367$
csbapr-no-nau (ReLU-MLP $+$ 6 fixes)	10	$40\%$	$0\%$	$0\%$	$0\%$	$-1{,}309$
csbapr (NAU $+$ 6 fixes)	10	$\mathbf{20\%}$	$10\%$	$10\%$	$0\%$	$-1{,}851$
csbapr-tanh (smooth-MLP $+$ 6 fixes)	10	$30\%$	$0\%$	$0\%$	$0\%$	$\mathbf{-1{,}257}$
csbapr-kan (KAN $+$ 6 fixes)	10	$30\%$	$10\%$	$0\%$	$0\%$	$-1{,}551$

Table 5: Crash rate at

m{=}20{\times}

across multiple severity thresholds. Each cell reports the fraction of seeds whose return at

m{=}20{\times}

falls below the threshold (lower is better). All six methods use

N_{\mathrm{seed}}{=}10

seeds. Thresholds show both typical poor performance and severe tail outcomes.

Figure 1: Static-demand return degradation relative to each configuration’s ID median. Every point is computed from Table 2 , with

N_{\mathrm{seed}}=10

for each underlying median. Lower degradation is better relative to that configuration’s own ID level. The figure reports a difference of aggregate medians and has no seed-level uncertainty bands.

Method	$m_{\mathrm{amp}}{=}1$	$m_{\mathrm{amp}}{=}2$	$m_{\mathrm{amp}}{=}5$	$m_{\mathrm{amp}}{=}10$	$m_{\mathrm{amp}}{=}20$	$m_{\mathrm{amp}}{=}50$
bapr (no SINDy, no NAU)	$\mathbf{-168}$	–	$-138$ K	$-1.28$ M	–	$\mathbf{-9.40}$ M
csbapr-no-sindy (NAU, no SINDy)	$-1{,}164$	–	$-105$ K	$-2.26$ M	–	$-11.85$ M
csbapr-no-nau (MLP $+$ SINDy)	$-2{,}301$	–	$-18$ K	$-0.94$ M	–	$-9.84$ M
csbapr (NAU $+$ SINDy, full)	$-2{,}916$	–	$\mathbf{-10}$ K	$\mathbf{-0.90}$ M	–	$-11.45$ M

Table 6: Exploratory LQR initial-amplitude sweep, median unaugmented episode return over

N_{\mathrm{seed}}=3

seeds. Higher is better. K and M denote

10^{3}

and

10^{6}

return units; – denotes an unreported cell. The symbolic-enabled configurations were labeled with SINDY_WITH_CONTROL=1 ; the implementation limitation is discussed in Section 7.1 . Bold : best reported value in the column.

Method	ID ( $\omega{=}0$ )	$\omega{=}1$	$\omega{=}2$	$\omega{=}5$	$\omega{=}10$	$\omega{=}20$
bapr-pristine (ReLU, no fixes)	$635$	$294$	$\mathbf{276}$	$85$	$127$	$197$
bapr / csbapr-relu (legacy, $N_{\mathrm{seed}}{=}2$ )	$995$	$216$	$221$	$144$	$99$	$107$
csbapr-tanh (smooth-MLP $+$ 6 fixes)	$975$	$\mathbf{330}$	$251$	$203$	$239$	$56$
csbapr-no-nau (ReLU-MLP $+$ 6 fixes)	$\mathbf{1{,}008}$	$317$	$270$	$\mathbf{255}$	$\mathbf{256}$	$\mathbf{224}$
csbapr (NAU $+$ 6 fixes)	$990$	$289$	$175$	$62$	$81$	$74$

Table 7: Hopper-v4 wind sweep, median episode return over

N_{\mathrm{seed}}=3

seeds except the legacy comparison (

N_{\mathrm{seed}}=2

, interrupted runs excluded). The pristine row disables the six stabilization settings; the legacy row is retained as a historical comparison. Wind amplitude

\omega=0

is ID. Bold : best in the column.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Quantity	Value
Discount $\gamma$	$0.99$
Target interpolation $\tau$	$0.005$
Initial entropy temperature $\alpha$	$0.2$ , automatically tuned
Optimizer step size $\eta_{\mathrm{opt}}$	$3\times 10^{-4}$
Critics $N_{\mathrm{crit}}$ ; batch size $N_{\mathrm{batch}}$	$5;\ 256$ ; transit critics: $3$
Run-length count $H_{\max}$ ; hazard $p_{\mathrm{haz}}$	$20;\ 0.02$

Appendix

Table 8: Shared defaults and the primary transit overrides.

Explore similar work

BAPR: Bayesian amnesic piecewise-robust reinforcement learning for non-stationary continuous control
May 15, 2026Yifan Zhang, Liang ZhengOffline Reinforcement Learning Bellman Equation
Central South University
AdamO: A Collapse-Suppressed Optimizer for Offline RL
May 3, 2026Nan Qiao, Sheng Yue, Shuning Wang +1Offline Reinforcement Learning Soft Actor-Critic
1Tsinghua University, China · 2Central South University, China · 3Sun Yat-sen University, China
Beyond Penalization: Diffusion-based Out-of-Distribution Detection and Selective Regularization in Offline Reinforcement Learning
May 6, 2026Qingjun Wang, Hongtu Zhou, Hang Yu +5Out-Of-Distribution Detection Out-Of-Distribution
School of Computer Science and Technology, Tongji University, Shanghai, China · MOE Key Lab of Embedded System and Service Computing, Tongji University, Shanghai, China · Shanghai Innovation Institute

May 15, 2026, cs.LGJ/K move · Enter open · S save

BAPR: Bayesian amnesic piecewise-robust reinforcement learning for non-stationary continuous control

Yifan Zhang, Liang Zheng

Central South University

Offline Reinforcement Learning Bellman Equation Change-Point Detection

Open paper AskPDF arXiv

Real-world control systems frequently operate under \emph{piecewise stationary} conditions, where dynamics remain stable for extended periods before undergoing abrupt regime changes. Standard robust RL methods face a fundamental dilemma: a globally conservative policy wastes performance during stable periods, while a locally adaptive policy risks catastrophic failure when the regime changes undetected. We propose \textbf{BAPR} (Bayesian Amnesic Piecewise-Robust SAC), which unifies Bayesian Online Change Detection (BOCD) with robust ensemble RL. The BAPR operator -- a convex combination of mode-conditional Bellman operators weighted by a frozen belief distribution -- is a

γ

-contraction. A complementary counterexample, machine-verified in Lean4, establishes a \emph{sharp boundary}: when beliefs depend on the Q-function, the contraction factor becomes $γ+ λΔ$ (where $Δ$ is the mode reward gap), and contraction fails exactly when $γ+ λΔ\geq 1$ . We derive a \emph{component-wise} formal error budget for the abstract operator -- every component machine-verified -- bounding post-switch recovery; the budget applies to the abstract mode-mixture operator and inherits to the implemented shared-critic algorithm only through the frozen-parameter design intuition. All results are formally verified with no \texttt{sorry} (1,145 lines across 3 Lean4 files, 22 machine-verified theorems). BOCD drives an adaptive conservatism mechanism: the policy becomes maximally conservative after detected change-points and smoothly relaxes as confidence grows, with detection delay

O(\log(1/δ))

. A context-conditioning module trained via RMDM loss provides mode-aware representations from simulator-provided mode IDs at training time and requires no mode labels at deployment.

Abstract

Explore similar work

BAPR: Bayesian amnesic piecewise-robust reinforcement learning for non-stationary continuous control

AdamO: A Collapse-Suppressed Optimizer for Offline RL

Beyond Penalization: Diffusion-based Out-of-Distribution Detection and Selective Regularization in Offline Reinforcement Learning