Towards Robust Reinforcement Learning for Small-Scale Language Model Agents

Md Rezwanul Haque; Md. Milon Islam; Fakhri Karray

cs.AIJul 27, 2026

Towards Robust Reinforcement Learning for Small-Scale Language Model Agents

Authors: Md Rezwanul Haque, Md. Milon Islam, Fakhri Karray

Abstract

The alignment of Small Language Models (SLMs) in the 70--500M parameter range using reinforcement learning is often considered unstable, though the underlying failure mechanisms have not been systematically investigated. In the State-of-the-Art (SOTA) research, fifteen (model, corpus) configurations were trained using Proximal Policy Optimization (PPO). The experiments included Pythia-70M, 160M, 410M and SmolLM2-135M, 360M on the TinyStories, CNN/DailyMail, and Wikitext-103 corpora. Three reproducible failure modes were identified in small-scale language models: silent LoRA parameter freezing in standard PEFT/TRL pipelines, numerical overflow in importance ratios when using bfloat16, and catastrophic policy collapse due to reward-model error. These issues were addressed using a merge-and-reinitialize adapter technique, float32 precision during PPO updates, and a three-layer safety mechanism comprising reward whitening, importance-ratio guarding, and weight rollback. In this paper, a capacity-headroom hypothesis is proposed, which states that PPO performance at the SLM scale depends on both a fluent supervised model ( $\text{PPL}<20$ ) and a discriminative reward signal, rather than on the number of model parameters. The proposed system converged stably in all experiments and improved preference win rate over the SFT baseline in configurations with a fluent prior and an informative reward signal. Furthermore, it outperformed instruction-tuned baselines while requiring significantly less training data. All checkpoints, preference datasets, and training scripts are publicly released $^§$ .

Explore similar work

Jun 8, 2026cs.AI

A Regret Minimization Framework on Preference Learning in Large Language Models

Reinforcement learning with verifiable rewards (RLVR) has enabled progress on reasoning-intensive tasks by relying on task-specific verifiers that provide automated correctness signals. However, many realistic language tasks are difficult to equip with reliable verifiers, motivating a growing reliance on reinforcement learning from human feedback (RLHF). In this setting, we argue that a closer examination of how human feedback should be interpreted is essential. We introduce Regret-based Preference Optimization

(\textbf{RePO})

, which reframes RLHF through

\textit{regret minimization}

rather than reward maximization. Human preferences are often shaped by

\textit{prospective}

anticipation of outcomes and

\textit{counterfactual}

comparisons to alternative behaviors, rather than by immediate, outcome-independent utility.

\textbf{RePO}

captures this structure by modeling preferences as behavior-conditioned assessments of relative suboptimality. Experiments on mathematical reasoning benchmarks and human preference datasets demonstrate consistent performance gains, indicating that

\textbf{RePO}

is an effective and human-aligned approach for training large language models.

Suhwan Kim, Taehyun Cho, Geon-Hyeong Kim +4

Reinforcement Learning From Human Feedback Preference Learning Verifiable Rewards Mathematical Reasoning Benchmarks Regret

Ask PDF

Jun 2, 2026cs.LG

RL Excursions during Pre-Training: Re-examining Policy Optimization for LLM training

The standard LLM training pipeline applies reinforcement learning (RL) only after pre-training and supervised fine-tuning (SFT). We question this status quo by training a LLM from scratch and applying RL, SFT, and SFT followed by RL directly to intermediate pre-training checkpoints. We find that RL is effective very early, and often matches the full SFT

\to

RL pipeline early as well. Through experiments on harder problems, we find that targeted pre-training data composition is a strong lever for RL effectiveness, even more so than model scale. Beyond reasoning accuracy, applying RL directly to base checkpoints expands the model's distribution; the sharpening effect reported in recent work arises only when RL follows SFT. The general capabilities of the model remain essentially unchanged by RL, while they degrade following SFT. Finally, we merge RL and SFT objectives by parallel averaging, which outperforms across all other training methods discussed, across metrics, while preserving general capabilities. Together, these results suggest that LLM training might benefit from an expanded use of RL.

Rachit Bansal, Clara Mohri, Tian Qin +2

Reinforcement Learning Post-Training Large Language Model Training Model Pretraining Safe Reinforcement Learning Model Checkpoints Policy Optimization

Ask PDF

Sep 17, 2026cs.LG

Score Centering Stabilizes Off-policy Reinforcement Learning

Reinforcement learning (RL) of large language models is notoriously sensitive to small differences between training and inference engines, often referred to as the training-inference mismatch (TIM). However, completely eliminating TIM is impractical, as it would come at a major cost to rollout efficiency. In this paper, we show that the instability of RL under TIM is primarily caused by drift: a persistent bias between training and inference engines that accumulates with every training step. We derive an additive "score centering" correction term that stabilizes RL under TIM by canceling drift. When training models from 0.6B to 30B parameters, score centering alone matches or outperforms methods based on importance sampling under quantization, with the gap growing as the mismatch becomes more severe. Because the correction is additive, score centering also composes with importance sampling -- their composition outperforms pure importance-sampling baselines in our staleness experiments.

Martin Marek, Max Ryabinin

Training-Inference Mismatch

Ask PDF