AIO1: an RL controller meant to work as a general APC optimizer (research checkpoints)

AIO1 is one recurrent soft actor-critic policy. It controls very different process plants through one plant-agnostic interface. A plant is described only by its public channels: controlled variables (CV) with setpoints, bands and scales, manipulated variables (MV) with ranges and slew limits, a CV–MV pairing map and the sample time. The policy is never given a plant ID, hidden states, true parameters or future disturbances. It is trained model-free on thousands of parallel simulated plants, with no teacher and no behaviour cloning of controllers.

Code, environments, reference controllers and the full research log are in the GitHub repository M3rcyzzz/AIOGym2. It is private for now, so this repo includes everything needed to load and run the weights. A live progress site is at m3rcyzzz.club/aio.

Status: research checkpoints, not a deployable controller. After fine-tuning, the policy beats or matches PI/DMC on some new plants. It loses to the Oracle on every plant. Fine-tuning made it forget two earlier tasks badly (see Results). Use it to study the method, not to control a real plant.

Checkpoints

Folder	What it is	Iteration / samples	Files
apc-v3-final-it30000/	Latest: final best.pt of the five-family fine-tune (Johansson quadruple tank, Shell heavy-oil fractionator, BSM1 activated sludge, Henson–Seborg pH, Tennessee Eastman), started from the mixed checkpoint below. The run ran its full 10 h.	it 30,000, 38.4 M samples, 10.0 h	actor.safetensors, full.safetensors, training.json, eval/
apc-v3-finetune/	The same fine-tune at its 5 h snapshot (best validation score at 5 h)	it 16,000, 20.48 M samples, 5.3 h	actor.safetensors, full.safetensors, training.json, eval/
mixed-it6000/	Best checkpoint of the four-tank + Tennessee Eastman mixed run	it 6,000, 7.68 M samples	actor.safetensors, full.safetensors, training.json, eval/

• actor.safetensors: the policy only (6.96 M parameters). Enough to run the controller.
• full.safetensors: actor.* plus both critics (critics.0.*, critics.1.*, 14.2 M parameters) and log_alpha. Use it to continue training.
• training.json: the checkpoint's iteration, samples, temperature, per-family reward scales, the validation row that selected it, and the full run configuration (command-line arguments, git commit).
• eval/: the per-scenario evaluation files the results below come from.

Which fine-tuned checkpoint to use depends on the plant. it 30,000 is better than it 16,000 on Shell and TEP, and worse on qtank, BSM1 and pH (see Evaluation). Both are kept because selection was noisy. it 30,000 was chosen on validation score, but one step later (it 30,001) its TEP validation IQM was 23,846, against 11,361 at it 30,000. Part of the margin that selected it is validation noise.

Network design

Code: modeling_aio.py, a standalone copy of experiments/memory_sac/network.py. About 21 M parameters in total: actor 7.0 M, two critics 7.1 M each.

1. Public token encoding (234 numbers). Up to 12 CV tokens × 12 features: error to setpoint, change, band edges relative to setpoint, tracking flag, measurement age, log sample time, setpoint change, presence mask. All features are divided by the CV's engineering scale. Up to 5 MV tokens × 6 features: position in range, last move, slew limit, mask. Then a 12 × 5 CV–MV pairing matrix. Every feature is dimensionless, so the same packet describes a four-tank rig, a distillation column or a wastewater plant. encode() in modeling_aio.py builds it.
2. Permutation-equivariant channel encoder. Shared per-token MLPs (symlog → Linear → LayerNorm → SiLU) embed the CV and MV tokens. Each MV token gathers the CV tokens paired with it, plus the pooled global CV embedding. The order of channels does not matter, and missing channels are masked out.
3. GRU memory (256). The pooled CV and MV summary goes through a GRU across the whole episode. The memory is how the policy identifies the plant from its own closed-loop history (gains, delays, disturbances). Nothing tells it which plant it is.
4. Gated mixture-of-experts head. Each MV token, concatenated with the memory, goes through a soft gate over 4 experts. Each expert is 3 pre-LayerNorm residual blocks of width 512. The actor outputs a mean and log-std per MV. The action is a latent z, and the plant applies u += tanh(z) · slew, clamped to the MV range, so a move never exceeds the slew limit.
5. Distributional critics. Each critic has its own encoder and GRU. The action is folded into the MV tokens, which are then pooled and fed to a MoE head. The head outputs 255 two-hot logits on a symlog support over [−20, 20]. The two critics are averaged, not minimised (FastTD3/FastSAC style).

Training method

• Learner: recurrent off-policy SAC in the style of FastSAC/FastTD3 (arXiv 2512.01996, 2505.22642).
  • Data: 1,024 GPU-batched plants, plus 256 Tennessee Eastman plants integrated in worker processes.
  • Batches: 256 sequences each, with a 16-step R2D2 burn-in and 16 trained steps; 4 gradient updates per vector step.
  • Optimiser: lr 3e-4, γ 0.99, τ 0.05, initial temperature 1e-3, target entropy −1 per active MV.
  • Failures: a physical failure is terminal, with reward −10⁴. A time-limit end is not bootstrapped.
• Scenario distributions. Every episode draws randomised plant parameters, operating points, setpoint changes and disturbances, so the policy has to adapt to the operating condition it is in.
• Reward: negative seconds outside the CV band, plus penalties for violating limits. Each plant family is divided by its own scale, measured from step cost in the warm-up. The scales are anchored on a family the loaded checkpoint already knows, so the critics' targets stay in range when training continues. Each family gets an equal share of every batch (BRC, arXiv 2505.23150).
• Model selection: geometric mean, over families, of the validation IQMs on fixed validation scenario sets. It is unit-free and needs no reference controller.
• Continuity: the model is never restarted from scratch. Each stage starts from the previous stage's weights, and every new family is first evaluated zero-shot. The chain:
  1. four-tank SAC;
  2. four-tank + TEP mixed run, best at it 6,000 (mixed-it6000/);
  3. five-family APC-suite v3 fine-tune: the 5 h snapshot (apc-v3-finetune/) and the final best (apc-v3-final-it30000/).
• Hardware: one RTX 4090 Laptop GPU. The mixed run took 8.65 h for 58.88 M samples; the fine-tune ran 10 h at about 1,060 samples/s (38.4 M samples).

Evaluation

Each plant has fixed test scenario sets (48 scenarios for qtank, Shell and BSM1, 60 for pH-P0). The policy runs deterministically. The metric is seconds outside the CV band plus seconds outside limits, per scenario; lower is better. Tables give the IQM over scenarios. The references are tuned PI (SIMC), tuned DMC, and an Oracle: an MPC/NMPC that knows the true model and the disturbances. The Oracle is a strong reference, not a strict lower bound.

Five-family fine-tune, test IQM

Plant	Zero-shot (mixed it 6,000)	Fine-tuned it 16,000 (5 h)	Fine-tuned it 30,000 (final)	PI	DMC	Oracle
Johansson quadruple tank	25,449	1,727	2,188	2,494	2,405	407
Shell fractionator	294,770	118,640	107,855	127,145	62,545	3,500
BSM1	2,353,013	155,662	159,188	493,275	376,200	34,050
pH (Henson–Seborg) P0	11,995 (13/60 failed)	409 (0 failed)	596 (0 failed)	671	194	93

Neither fine-tuned checkpoint fails on these sets. In the bullets below, the paired comparisons (win counts and 95% bootstrap CIs) are for it 16,000. The ones for it 30,000 are in its eval/ folder. One more point about it 30,000 on Shell: its IQM is lower, but its constraint-violation time is about 4× that of it 16,000 (119,400 s against 28,920 s).

• Zero-shot transfer did not happen. On the three new plants, zero-shot was worse than holding the MVs still (hold IQM: qtank 12,615, Shell 283,330, BSM1 2,037,113).
• After fine-tuning:
  • qtank: level with PI and DMC. It wins 32/48 against PI and 30/48 against DMC, but the 95% CIs of the paired IQM difference cross 0.
  • Shell: a little better than PI (27/48, CI [−21,715, 0]). It clearly loses to DMC (5/48).
  • BSM1: beats PI and DMC in 48 of 48 paired scenarios.
  • pH: lower IQM than PI (42/60 wins, CI [−388, +101], not significant). It loses to DMC (9/2/49).
• The Oracle beats the policy on every plant.

Forgetting (the important negative result)

Earlier task	Mixed it 6,000	Fine-tuned it 16,000	Fine-tuned it 30,000
Old four-tank P0 v2, IQM	728 (better than DMC 40/0/8, 47.5 s worse than Oracle)	7,129 (31 of 48 scenarios fail physically)	7,571 (22 of 48 fail physically)
Tennessee Eastman TEP-P0, IQM	9,244 (level with DMC; worse than setpoint PI by 681 s and than the Oracle)	17,815 (loses 46 of 48 to setpoint PI and DMC)	12,934 (setpoint PI 8,683, DMC 9,204)

Between it 16,000 and it 30,000, TEP partly recovered and the old four-tank did not. TEP is one of the five fine-tuning families; the old four-tank is not.

The selection score is a geometric mean over the families being trained. Progress on the new families hid the regression on the old ones. The repo now has the forgetting guards from Wolczyk et al. (ICML 2024): behaviour cloning toward the earlier policy on the old families, plus feasibility-first checkpoint selection with a regression limit (rsac --anchor/--guard-families). A run with these guards (apc-v3/anchor-seed1) started when the fine-tune ended. Its results will be added here when it finishes.

Mixed run (mixed-it6000/)

Four-tank + TEP, at it 6,000 (7.68 M samples):

• Old four-tank P0 v2: IQM 728. Better than DMC (40 wins, 0 ties, 8 losses), 47.5 s worse than the Oracle.
• TEP-P0: IQM 9,244. Level with DMC; worse than setpoint PI and the Oracle. Constraints are handled better than DMC, and tracking is worse.
• Later in that run TEP validation degraded while the four-tank part improved. Selection on the geometric mean kept the it 6,000 checkpoint.

Usage

# pip install torch numpy safetensors huggingface_hub
from huggingface_hub import snapshot_download
import sys
path = snapshot_download("M3rcyzzz/AIO1", allow_patterns=["modeling_aio.py", "apc-v3-final-it30000/actor.safetensors"])
sys.path.insert(0, path)
import modeling_aio as M

actor = M.load_actor(f"{path}/apc-v3-final-it30000/actor.safetensors")
ctl = M.Controller(actor)            # keeps the GRU state; call ctl.reset() at a new episode

# every decision step, from your plant's public signals (numpy arrays, one entry per CV / MV):
packet = M.encode(values=y, previous=y_prev, target=sp, last_target=sp_prev, scale=cv_scale,
                  lower=cv_low, upper=cv_high, tracking=is_tracked, command=u, last_command=u_prev,
                  u_low=u_min, u_high=u_max, slew=du_max, pairing=pairing, dt=dt)
move = ctl.step(packet)              # tanh(z) per MV, in [-1, 1]
u = np.clip(u + move * du_max, u_min, u_max)

To continue training, load full.safetensors with M.load_full (it returns the actor, the critics and log_alpha). Training code is experiments/scenario_ppo/rsac.py in the (currently private) GitHub repo. export.py converts an rsac .pt checkpoint into this layout and checks that the actor outputs match exactly.

Limitations

• Every plant is a simulation (pc-gym-derived plants, the Johansson quadruple tank, the Shell fractionator, BSM1, Henson–Seborg pH, the native Tennessee Eastman code). Nothing has been validated on a real plant.
• The policy is not safety-certified. In production it would need a supervisory or fallback layer.
• There is no zero-shot transfer to unseen plant families yet, and fine-tuning causes forgetting (see above).

License

MIT, the same as the GitHub repository (Copyright (c) 2026 SUPCON International).

---

中文简介

AIO1 是一个用模型无关强化学习训练的循环 SAC 控制器，目标是做成“通用 APC 优化器”。

• 接口：一个装置只用公开通道描述，即 CV/SP/带宽/量程、MV 范围与限速、CV–MV 配对和采样时间。不输入装置 ID、隐藏状态、真参数或未来扰动。
• 网络：234 维无量纲 token，经置换等变通道编码器、256 维 GRU 记忆、4 专家门控残差头，以及 255 原子 two-hot 分布式双 critic。
• 训练：上千个并行仿真装置，每回合随机工况。各族奖励按尺度归一化，按验证集 IQM 的几何平均选模。每个阶段都从上一阶段的权重接着训。

结果如实写出：

• 新装置零样本没有迁移。
• 五族微调 5 h 快照（it 16000）：四容水箱与 PI/DMC 持平；Shell 略好于 PI（不显著）、明显输 DMC；BSM1 48 个场景全胜 PI 和 DMC；pH 的 IQM 低于 PI（不显著），输 DMC。
• 跑满 10 h 的最终 best（it 30000）：Shell 和 TEP 比 5 h 快照好，四容水箱、BSM1、pH 变差；Shell 越限时间约为 5 h 快照的 4 倍。选模依据的验证分数噪声很大，两份权重都保留。
• 所有装置都输 Oracle。
• 旧四水箱遗忘严重，到 it 30000 也没有恢复（P0 v2 IQM 728→7,571，22/48 物理失败）。TEP 先退化，后来部分恢复（9,244→17,815→12,934）。带防遗忘约束（行为克隆锚和选模退化门槛）的一轮已经开跑，结果出来后补上。