Pick-to-Learn Safety Fine-Tuning
Implementation & benchmark of arXiv:2607.21086 — Certified Stochastic Control via Covariance Steering with Pick-to-Learn
Source paper: Certified Stochastic Control via Covariance Steering with Pick-to-Learn arXiv:2607.21086 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Train a neural policy against a simulator using an adaptive constraint set formed from the worst violations, rather than uniformly averaging all rollouts. At each round, identify the trajectory with the largest normalized safety violation, add its state-time features and violation margin to a surrogate barrier or penalty model, and fine-tune the policy until the surrogate constraints are satisfied. This should reduce the gap between nominal validation risk and rare-event failure risk while concentrating simulator calls on informative trajectories.
Formulas
Mathematical statement
A rollout samples initial state $x_0\sim\mu_0$, fixed uncertain parameter $\lambda\sim\mu_\Lambda$, and Brownian noise $w_t$, then evolves according to $dx_t=f(x_t,u_t,t;\lambda)dt+Gdw_t$, where $x_t$ is the state, $u_t=\pi_\theta(x_t,t)$ is the neural policy, $f$ is the drift, $G$ is the diffusion matrix, and $\theta$ are policy parameters. Let $g_r(\tau)\le 0$ be the $r$-th safety constraint evaluated on trajectory $\tau$, and define normalized violation $v(\tau)=\max_r[g_r(\tau)/s_r]_+$ with scale $s_r>0$ and $[z]_+=\max(z,0)$. Pick-to-Learn adds the maximizer $\tau^*=\arg\max_{\tau\in\mathcal{D}_m}v(\tau)$ from the current simulator batch to the surrogate constraint set $\mathcal{C}_{m+1}=\mathcal{C}_m\cup\{(\tau^*,v(\tau^*))\}$. A conservative empirical certificate uses $\widehat p_v=N_v/N$ and an upper confidence correction $\epsilon(N,\delta,\mathcal{K})$, giving $p_v\le\widehat p_v+\epsilon$ with confidence at least $1-\delta$; $N_v$ is the number of violating test rollouts, $N$ is the test-set size, and $\mathcal{K}$ is the effective compressed policy/constraint description size. The mechanism predicts that worst-case violation decreases sharply after repeatedly adding high-$v$ trajectories, while the certified upper bound decreases approximately as $O(\sqrt{(\mathcal{K}+\log(1/\delta))/N})$ when the policy is held fixed during certification.
Implementation notes
1. Integration point: use this as an outer loop around an RL or differentiable-simulator policy update. Keep the neural policy architecture and task loss unchanged, but add a replay buffer of simulator trajectories and a differentiable safety loss. For each constraint $g_r$, normalize by a fixed engineering scale $s_r$ so unrelated units do not dominate. 2. Pseudocode: initialize policy parameters $\theta$, constraint buffer $\mathcal{C}=\varnothing$, and repeat for $m=1,\ldots,M$: sample $B$ trajectories with independently sampled $x_0$, fixed-per-trajectory $\lambda$, and process noise; compute $v(\tau)$ for every trajectory; select the top-$q$ trajectories, or at minimum $\tau^*=\arg\max v$; append them to $\mathcal{C}$; perform $E$ gradient steps on $L_{task}(\theta)+\beta|\mathcal{C}|^{-1}\sum_{\tau\in\mathcal{C}}\operatorname{softplus}(v_\theta(\tau)/T)^2$, where $T$ is a smoothing temperature. Stop adding constraints when the batch maximum is below tolerance. 3. Computed from the paper: simulator dynamics, fixed uncertain parameters, stochastic rollouts, worst-violation selection, and the compression-style risk certificate. Estimated empirically: the policy-dependent gradients, effective description size $\mathcal{K}$, and violation probability. Use a held-out certification set generated only after training; do not reuse adaptive trajectories for $\widehat p_v$. 4. First cheap experiment: train a small MLP controller on a two-dimensional stochastic point-mass reaching task with random mass or gravity, comparing uniform-risk training, nominal covariance training, and Pick-to-Learn with equal simulator budget. Use 600 held-out rollouts and confidence $1-\delta=0.95$. The quantitative prediction is that the maximum observed normalized violation and held-out violation rate drop more rapidly per simulator rollout under Pick-to-Learn, while the certified bound follows an $N^{-1/2}$ slope when plotted against test-set size. A falsifiable target is a factor-of-two reduction in held-out violations at equal budget and agreement between the empirical rate and the upper bound within the expected confidence slack.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: Significantly worse ✗
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built a self-contained Pick-to-Learn safety fine-tuning toy experiment with worst-violation selection, differentiable safety penalties, held-out certification, and fixed seeds. The mechanism manifested: observed order-statistic maxima matched B/(B+1) within 0.001, the certificate correction had log-log slope -0.500 versus the predicted -0.5, and Pick-to-Learn reduced held-out violation rate from 0.2331 to 0.0837 at equal simulator budget.
- Agent confidence
- 8/10
- Baseline
- Uniform updates: theta=0.8085, violation rate=0.2331, maximum normalized violation=0.1885, mean violation=0.02197.
- Idea
- Pick-to-Learn: theta=0.6766, violation rate=0.0837, maximum normalized violation=0.0566, mean violation=0.002368. Held-out N=600: 52 violations, phat=0.0867, 95% Hoeffding upper bound=0.1366. Certificate epsilon slope=-0.500.
Limitations: This is a one-step scalar stochastic safety abstraction rather than a 2D multi-step simulator or full neural policy. It does not test RL-scale sample efficiency, Brownian dynamics, top-q selection, model misspecification, or a K-dependent certificate.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.21086)
Significantly WORSE than the tuned baseline (+373376.2%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0
- Idea mean
- 0.1065
- Effect (Δ)
- 0.1065 (+373376.2%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±60085.4%
- Mechanism
- Not confirmed ✗
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
Only the registered built-in dynamics supervised pendulum benchmark was tested. The implementation uses predicted angle as a surrogate safety signal rather than a true multi-step closed-loop policy constraint, and the fixed held-out dataset cannot measure policy-induced changes in physical violation frequency.
How to run: python3 bench_stage2.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 10/10
Pick-to-Learn safety fine-tuning was evaluated on the registered built-in dynamics track with the matched rnn_small architecture. The complete protocol ran 8 paired seeds, baseline tuning over the shared learning-rate grid, and an equal-size idea sweep. The idea reduced its predicted safety violation rate by 62%, but significantly worsened the independent test MSE, so the benchmark verdict is idea worse (significant) and worked=false.
- Baseline
- Uniform MSE baseline, best lr=0.01: mean test MSE 2.8511603204606217e-05 over 8 seeds.
- Idea
- Top-8 normalized predicted-angle violation replay, best lr=0.01: mean test MSE 0.10648404061794281 over 8 seeds; predicted violation rate 0.15875 versus baseline 0.4179166666666667; observed held-out violation rate 0.41708333333333336.
Benchmark result
Significantly WORSE than the tuned baseline (+373376.2%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Dynamics forecasting
- Model
- rnn_small
- Paired seeds
- 8
- Baseline mean
- 0
- Idea mean
- 0.1065
- Effect (Δ)
- 0.1065 (+373376.2%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±60085.4%
- Mechanism
- Not confirmed ✗
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
Only the registered built-in dynamics supervised pendulum benchmark was tested. The implementation uses predicted angle as a surrogate safety signal rather than a true multi-step closed-loop policy constraint, and the fixed held-out dataset cannot measure policy-induced changes in physical violation frequency.
How to run: python3 bench_stage2.py
Artifacts
- bench_report.json 6.1 KB View Raw
- bench_stage2.py 7.3 KB View Raw
- experiment.py 5.1 KB View Raw
- report.md 1.3 KB View
- report_bench_2026-08-31T155054.md 8.1 KB View
- results.json 9.1 KB View Raw