ICML 2026 Papers

International Conference on Machine Learning

Decision breakdown

6,341papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Spotlight
    5368.5%
  • Poster
    5,80591.5%

Browse all papers

Topic: Reinforcement Learning · 1,021 papers · Reset filters

Reward-Preserving Counterfactual State Editing for Offline Reinforcement Learning

Siyu Wang, Xiaocong Chen, Mingming Gong, Yong Li, Quan Z. Sheng, Lina Yao

Expand abstract

Transformer sequence models such as Decision Transformer can learn strong offline policies from logged trajectories, but they often suffer from causal confusion: reliance on spurious correlations that predict reward in the data but do not reflect the true causal mechanisms of the environment. We propose CSET (Counterfactual State Editing Transformer), which improves robustness in strictly offline reinforcement learning without learning environment transition dynamics. On the data side, CSET fits a causal reward model as a conditional variational autoencoder and uses a counterfactual state generator to propose minimally edited observations whose predicted reward matches the factual reward, under a normalized move-band constraint and an acceptance gate that enforce plausibility and reward consistency; augmentation replaces only the observation token to avoid synthetic successor transitions. On the model side, CSET uses a causally structured hybrid transformer: modality-specific convolutional encoders process return-to-go, state, and action streams, and a final attention block is softly supervised so action prediction focuses on its direct causal parents. Experiments on D4RL locomotion, AntMaze, and offline recommendation benchmarks show consistent gains within the DT family, and CSET remains substantially more robust than strong value-based and DT baselines under injected spurious distractors.

Poster
From Teacher Pathways to Invariant Manifolds: Consensus Subspace Distillation for TSFMs

Zexing Zhang, Tianyang Lei, Jichao Li, Yang Kewei

Expand abstract

Time-series foundation models (TSFMs) deliver strong cross-domain generalization, but their scale makes deployment costly. Knowledge distillation is a natural compression route, yet prior TSFM distillation typically imitates teacher outputs, features, or pairwise relations, and therefore remains tightly coupled to teacher-specific training trajectories while underutilizing two empirical properties: (i) high-level representations across model scales tend to converge toward a shared, approximately low-rank geometry, and (ii) layer-wise utility follows a long-tail pattern. We propose consensus subspace distillation, which reframes distillation as aligning a student to a model-agnostic geometric object: a scale-invariant low-rank consensus subspace together with its center statistics. Offline, we screen high-contribution layers via drop-layer marginal loss, estimate a shrinkage-stabilized covariance from their embeddings, and derive a truncated eigensubspace that defines a consensus projector. Online, we project student embeddings into this subspace and match the teacher’s projected mean and covariance using a lightweight mean--covariance objective, enabling stable optimization without rigid pointwise feature binding. To mitigate subset-induced bias, we further introduce a frequency-domain uncertainty injection mechanism that inflates spectral density based on characteristic-function discrepancies and injects dispersion only within the consensus directions. Across forecasting and imputation, the distilled student matches or slightly improves upon the teacher, while exhibiting a predictable trade-off under strict zero-shot classification. With MOMENT-Large as teacher, we achieve about 90% parameter reduction and substantial distillation-time savings while retaining comparable performance across multiple time-series tasks. Code and compressed weights are available at anonymous.4open.science/r/CSD-13C3/.

Poster
Are Tools Always Beneficial? Learning to Invoke Tools Adaptively for Dual-Mode Multimodal LLM Reasoning

Qinghe Ma, Zhen Zhao, Yiming Wu, Jian Zhang, Lei Bai, Yinghuan Shi

Expand abstract

Tool-augmented reasoning has emerged as a promising direction for enhancing the reasoning capabilities of multimodal large language models (MLLMs). However, existing studies mainly focus on enabling models to perform tool invocation, while neglecting the necessity of invoking tools. We argue that tool usage is not always beneficial, as redundant or inappropriate invocations largely increase reasoning overhead and even mislead model predictions. To address this issue, we introduce AutoTool, a model that adaptively decides whether to invoke tools according to the characteristics of each query. Within a reinforcement learning framework, we design an explicit dual-mode reasoning strategy with mode-specific reward functions to guide the model toward producing accurate responses. Moreover, to prevent premature bias toward a single reasoning mode, AutoTool jointly explores and balances tool-assisted and text-centric reasoning throughout training, and promotes free exploration in later stages. Extensive experiments demonstrate that AutoTool exhibits outstanding performance and high efficiency, yielding a 21.8\% accuracy gain on V* benchmark compared to the base model, and a 44.9\% improvement in efficiency over existing tool-augmented methods on POPE benchmark. Code is available at https://github.com/MQinghe/AutoTool.

Poster
Return-Critic: Bridging Goal Discrepancy for Efficient Visual Reinforcement Learning

Ruyi Lu, Xuesong Wang, Hengrui Zhang, Yuhu Cheng

Expand abstract

Sample inefficiency remains a challenge in pixel-based visual reinforcement learning (RL), primarily due to ineffective state representation learning. While recent advances employ auxiliary tasks to improve representation learning, their representation goals (e.g., mask reconstruction, state prediction) are misaligned with the ultimate RL goal of maximizing return, constraining further improvements in representation quality. To achieve efficient visual reinforcement learning, we propose Return-Critic (RC), an auxiliary framework that bridges goal discrepancy by return prediction. RC samples partial frames from an episode, processes them through a shared visual encoder, and employs a lightweight Transformer to predict the episode's return, forcing the encoder to learn return-relevant representation. The attention weights naturally highlight important frames, enabling a key function for prioritized learning. Extensive experiments on both online (DMControl) and offline (V-D4RL) benchmarks demonstrate that RC significantly enhances the sample efficiency, particularly achieving 68% performance boost on average across nine challenging tasks from DMControl.

Poster
Near-Optimal Regret for Policy Optimization in Contextual MDPs with General Offline Function Approximation

Orin Levy, Aviv Rosenberg, Alon Cohen, Yishay Mansour

Expand abstract

We introduce OPO-CMDP, the first policy optimization algorithm for stochastic Contextual Markov Decision Process (CMDPs) under general offline function approximation. Our approach achieves a high probability regret bound of $\widetilde{O}(H^4\sqrt{T|S||A|\log(|\mathcal{F}||\mathcal{P}|)}),$ where $S$ and $A$ denote the state and action spaces, $H$ the horizon length, $T$ the number of episodes, and $\mathcal{F}, \mathcal{P}$ the finite function classes used to approximate the losses and dynamics, respectively. This is the first regret bound with optimal dependence on $|S|$ and $|A|$, directly improving the current state-of-the-art (Qian, Hu, and Simchi-Levi, 2024). These results demonstrate that optimistic policy optimization provides a natural, computationally superior and theoretically near-optimal path for solving CMDPs.

Poster
Scaling the Scaling Logic: Agentic Meta-Synthesis of Logic Reasoning

Bowen Liu, Zhi Wu, Runquan Xie, Zhanhui Kang, Jia Li

Expand abstract

Reinforcement Learning from Verifiable Rewards (RLVR) is bottlenecked by data: existing synthesis pipelines rely on expert-written code or fixed templates, confining growth to instance-level perturbations. We shift the evolvable unit from problem instances to task-family specifications. SSLogic is an agentic meta-synthesis framework in which LLM agents iteratively author and refine executable Generator-Validator pairs inside a closed Generate-Validate-Refine loop, producing families with new rules and difficulty gradients rather than parameter variations of old ones. A Multi-Gate Validation Protocol, multi-strategy consensus plus Adversarial Blind Review, where independent agents solve each instance by writing and executing code, filters ill-posed tasks before they enter training. Starting from 400 seed families, two evolution rounds yield 953 families and 21,389 verifiable instances. Three converging comparisons (step-matched, token-matched, and size-controlled on external Enigmata data) consistently show higher training utility of evolved data, with gains of SynLogic +5.2, AIME25 +3.0, and BBH +5.5 on Enigmata. Fine-grained KORBench evaluation reveals selective improvements in logic (+13.2%) and operation (+9.6%), linking structural evolution to downstream gains. Code is available at https://github.com/AdAstraAbyssoque/Scaling-the-Scaling-Logic.

Poster
Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatch

ziyang zhang, Xinheng Ding, Jiayi Yuan, Rixin Liu, Huizi Mao, Jiarong Xing, Zirui Liu

Expand abstract

Deterministic inference is increasingly critical for Large language model (LLM) applications such as LLM-as-a-judge evaluation, multi-agent systems, and reinforcement learning (RL). However, existing LLM serving frameworks can produce different outputs for identical inputs when tensor parallel (TP) size or batch size changes, even under greedy decoding. This arises from the non-associativity of floating-point arithmetic and inconsistent reduction orders across GPUs. While prior work has addressed batch-size–related nondeterminism through batch-invariant kernels, determinism across different TP sizes remains an open problem, particularly in RL settings, where the training engine typically uses Fully Sharded Data Parallel (FSDP) (i.e., TP = 1) while the rollout engine relies on multi-GPU TP to maximize the inference throughput, creating a probability mismatch that can degrade or even collapse training. We identify and analyze the root causes of TP-induced inconsistency and propose **Tree-Based Invariant Kernels (TBIK)**, a set of custom matrix multiplication and reduction kernels that guarantee bit-wise identical results across TP sizes. Our key insight is to enforce a consistent reduction order across and within GPUs. We implement TBIK in Triton and integrate it into vLLM and FSDP, achieving **bit-wise deterministic inference** across different TP sizes and **zero probability divergence** between rollout and training engines in RL pipelines. By eliminating mismatches caused by different parallelization strategies, TBIK enables **true on-policy RL at scale for the first time**, leading to improved model performance and faster convergence.

Poster
PAC-Bayesian Reinforcement Learning Trains Generalizable Policies

Abdelkrim ZITOUNI, Mehdi hennequin, Juba Agoun, Ryan Horache, NADIA KABACHI, Omar Rivasplata

Expand abstract

We derive a novel PAC-Bayesian generalization bound for reinforcement learning that explicitly accounts for Markov dependencies in the data, through the chain's mixing time. This contributes to overcoming challenges in obtaining generalization guarantees for reinforcement learning, where the sequential nature of data breaks the independence assumptions underlying classical bounds. The new bound provides non-vacuous certificates for modern off-policy algorithms such as Soft Actor-Critic. We demonstrate the practical utility of the bound through PB-SAC, a novel algorithm that optimizes the bound during training to guide exploration. Experiments across several continuous control tasks show that the proposed approach provides meaningful confidence certificates while maintaining competitive performance.

Poster
DeepHA: Scaling Action Chains Elicits Deep Hierarchical Agents

Zihao Wang, Muyao Li, Kaichen He, Haowei Lin, Xiaojian Ma, Anji Liu, Yitao Liang

Expand abstract

Prevailing autonomous agents are often constrained by a single, predefined action space, which limits their generalization capabilities across diverse tasks and can introduce compounding errors through decoupled policy execution. To address these limitations, we introduce the Deep Hierarchical Agent (DeepHA), a unified architecture that operates across a mixture of heterogeneous action spaces, flexibly generating actions ranging from high-level semantic skills to low-level motor controls. We further propose a Chain-of-Action (CoA) reasoning framework, which enables the agent to use higher-level abstract actions as structured `thoughts' to guide the generation of more granular, subsequent actions. To manage the computational demands of this deep reasoning in long-horizon tasks, we develop a memory-efficient mechanism that dynamically compresses historical context and leverages Key-Value (KV) caching, reducing context length by approximately 75% without sacrificing performance. We conduct extensive evaluations on a new, large-scale benchmark of over 800 diverse Minecraft tasks. Results show that DHA significantly outperforms prior methods, establishing a new state-of-the-art and demonstrating superior generalization, particularly in complex, multi-step planning tasks. Our work presents a novel, unified framework for building more capable and efficient autonomous agents.

Poster
UI2Code^N: UI-to-Code Generation as Interactive Visual Optimization

Zhen Yang, Wenyi Hong, Mingde Xu, Xinyue Fan, Weihan Wang, Jiale Cheng, Xiaotao Gu, Jie Tang

Expand abstract

UI-to-code aims to translate UI screenshots into executable front-end code. Despite progress with vision-language models (VLMs), most existing methods formulate UI-to-code as a single-pass generation, which mismatches real-world UI development that is inherently iterative and feedback-driven. We reformulate UI-to-code as an interactive visual optimization problem, where code generation is embedded in a closed-loop process of execution, visual inspection, and iterative refinement driven by rendered visual feedback. To address the non-differentiability of visual objectives and the noise of absolute visual evaluators, we propose Relative Visual Policy Optimization (RVPO), a preference-based reinforcement learning method that optimizes relative visual rankings among rendered candidates under execution feedback. We instantiate this paradigm in UI2Code$^{\text{N}}$, an open-source 9B model trained via continual pre-training, supervised fine-tuning, and reinforcement learning. Experiments demonstrate state-of-the-art performance on UI drafting, UI polishing, and UI editing benchmarks, even outperforming larger models, with performance consistently improving through iterative visual optimization. Our code and models are available at \url{https://github.com/zai-org/UI2Code_N}.

Poster
DREAM-R: Multimodal Speculative Reasoning with RL-Based Refined Drafting, Precise Verification, and Fully Parallel Execution

Yunhai Hu, Zining Liu, Xiangyang Yin, Tianhua Xia, BO BAO, Eric Sather, Vithursan Thangarasa, Sai Qian Zhang

Expand abstract

Speculative reasoning has recently been proposed as a means to accelerate reasoning-intensive generation in large multimodal models, but its effectiveness is often constrained by misalignment between speculative drafts and target-verified reasoning. In this work, we introduce \textit{DREAM-R}, a framework that substantially improves the performance of speculative reasoning. At its core, DREAM-R employs \textit{Speculative Alignment Policy Optimization} (SAPO), a reinforcement-learning objective that trains draft models to generate reasoning steps that are both faithful to target trajectories and concise. We further propose a \textit{Threshold-based Verification Mechanism} (TBVM) that uses a ratio-based criterion to provide stable and interpretable acceptance of speculative steps only when positive evidence clearly dominates, thereby preventing error propagation. Building on these components, we develop a \textit{Fully Parallel Speculative Reasoning} (FPSR) framework that parallelizes draft generation, target-side reasoning, and verification across multi-step reasoning, enabling early stopping and clean fallback. Experiments on reasoning-heavy benchmarks demonstrate up to $2.49\times$ speedup while preserving target-model accuracy, yielding substantial efficiency gains without compromising reasoning quality.

Poster
UAV$^2$: A Unified and Adaptive Scheduling Framework for UAV Autopilot System with Reinforcement Learning

Zeying Li, Shuai Zhao, Chaowen Wu, Boyang Li, Kai Huang

Expand abstract

Unmanned aerial vehicle (UAV) autopilot systems typically comprise navigation and flight-control modules, and their effective scheduling is critical to achieving high flight performance. However, most existing UAV platforms adopt a split architecture in which navigation and flight control are deployed on separate hardware devices. This separation restricts system-wide observability and prevents holistic scheduling and optimization across the entire autopilot pipeline. Moreover, autonomous flight performance emerges from implicit, cross-coupled, and accumulated interactions among multiple factors, rendering traditional model-based or heuristic scheduling approaches ineffective. To address these challenges, we propose UAV$^2$, a unified and adaptive scheduling framework for UAV autopilot systems with reinforcement learning, targeting flight performance optimization. UAV$^2$ integrates navigation and flight control onto a single onboard computing platform and operating system, formulates the scheduling problem as a partially observable Markov decision process, and learns scheduling policies from runtime execution feedback. The proposed approach is trained and evaluated in a hardware-in-the-loop simulation environment. Experimental results demonstrate that the learned scheduling policy consistently outperforms fixed-rate scheduling strategies in terms of flight robustness and tracking performance.

Poster
Thinking in Latent Space: Progressive Multimodal Simplification for Visual Reasoning

Yuesen Tang, Yiming Yang, Tengfei Bao, Yu Tong

Expand abstract

Recent Multimodal Large Language Models (MLLMs) have advanced cross-modal reasoning by extending Chain-of-Thought (CoT) prompting to visual tasks. However, existing methods still rely heavily on explicit textual reasoning steps, leading to information loss, unstable perception–reasoning interaction, and high computational cost. Inspired by human cognition, we argue that effective visual reasoning emerges from a dynamic interplay between perception and latent thought, rather than a purely linear verbalization process. Motivated by this insight, we propose Latent-Driven Progressive Visual Reasoning (LDPVR), a framework that formulates multimodal reasoning as a Markov Chain of Recursive State Simplification, where explicit textual states are progressively refined under the guidance of latent transitions. Central to LDPVR is Interleaved Latent Grounding, which leverages latent semantic intent to actively retrieve fine-grained visual evidence and drive robust state evolution, enabling the model to iteratively reduce uncertainty before committing to simplified textual states. To optimize this process, we introduce a three-stage curriculum combining supervised fine-tuning, latent-text distillation, and reinforcement learning via Group Relative Policy Optimization (GRPO). Experiments on six multimodal reasoning benchmarks demonstrate that LDPVR improves reasoning accuracy while maintaining low inference latency.

Poster
RAD: Retrieval High-quality Demonstrations to Enhance Decision-making

Lu Guo, Yixiang Shan, Zhengbang Zhu, Qifan Liang, Lichang Song, Ting Long, Weinan Zhang, Yi Chang

Expand abstract

Offline reinforcement learning (RL) learns policies from fixed datasets, thereby avoiding costly or unsafe environment interactions. However, its reliance on finite static datasets inherently restricts the ability to generalize beyond the training distribution. Prior solutions based on synthetic data augmentation often fail to generalize to unseen scenarios in the (augmented) dataset. To address these challenges, we propose Retrieval High-quAlity Demonstrations (RAD) for decision-making, which innovatively introduces a retrieval mechanism into offline RL. Specifically, RAD retrieves high-return and reachable states from the offline dataset as target states, and leverages a generative model to generate sub-trajectories conditioned on these targets for planning. Since the targets are high-return states, once the agent reaches such a target, it can continue to obtain high returns by following the associated high-return actions, thereby improving policy generalization. Extensive experiments confirm that RAD achieves competitive or superior performance compared to baselines across diverse benchmarks, validating its effectiveness. Our code is available at https://github.com/LeahGL/RAD.

Poster
Learning Query-Aware Budget-Tier Routing for Runtime Agent Memory

Haozhen Zhang, Haodong Yue, Tao Feng, Quanyu Long, Jianzhu Bao, Bowen Jin, Weizhi Zhang, Xiao Li, Jiaxuan You, Chengwei Qin, Wenya Wang

Expand abstract

Memory is increasingly central to Large Language Model (LLM) agents operating beyond a single context window, yet most existing systems rely on offline, query-agnostic memory construction that can be inefficient and may discard query-critical information. Although runtime memory utilization is a natural alternative, prior work often incurs substantial overhead and offers limited explicit control over the performance-cost trade-off. In this work, we present **BudgetMem**, a runtime agent memory framework for explicit, query-aware performance–cost control. BudgetMem structures memory processing as a set of memory modules, each offered in three budget tiers (i.e., Low/Mid/High). A lightweight router performs budget-tier routing across modules to balance task performance and memory construction cost, which is implemented as a compact neural policy trained with reinforcement learning. Using BudgetMem as a unified testbed, we study three complementary strategies for realizing budget tiers: implementation (method complexity), reasoning (inference behavior), and capacity (module model size). Across LoCoMo, LongMemEval, and HotpotQA, BudgetMem surpasses strong baselines when performance is prioritized (i.e., high-budget setting), and delivers better accuracy–cost frontiers under tighter budgets. Moreover, our analysis disentangles the strengths and weaknesses of different tiering strategies, clarifying when each axis delivers the most favorable trade-offs under varying budget regimes. Code is available at https://github.com/ViktorAxelsen/BudgetMem

Poster
AdvEvo-MARL: Shaping Internalized Safety through Adversarial Co-Evolution in Multi-Agent Reinforcement Learning

Zhenyu Pan, Yiting Zhang, Zhuo Liu, Yolo Y. Tang, Zeliang Zhang, Haozheng Luo, Chenwei Xu, Yuwei Han, Jianshu Zhang, Dennis Wu, Hong-Yu Chen, Haoran Lu, Haoyang Fang, Manling Li, Chenliang Xu, Philip S. Yu, Han Liu

Expand abstract

LLM-based multi-agent systems excel at planning, tool use, and role coordination, but their openness and interaction complexity also expose them to jailbreak and adversarial collaboration. Existing defenses fall into two lines: (i) self-verification that asks each agent to pre-filter unsafe instructions before execution, and (ii) external guard modules that police behaviors. The former often underperforms because a standalone agent lacks sufficient capacity to detect cross-agent unsafe chains and delegation-induced risks; the latter increases system overhead and creates a single-point-of-failure—once compromised, system-wide safety collapses, and adding more guards worsens cost and complexity. To solve these challenges, we propose AdvEvo-MARL, a co-evolutionary multi-agent reinforcement learning framework that internalizes safety into task agents. Rather than relying on external guards, AdvEvo-MARL jointly optimizes attackers (which synthesize evolving jailbreak prompts) and defenders (task agents trained to both accomplish their duties and resist attacks) in adversarial learning environments. To stabilize learning and foster cooperation, we introduce a public baseline for advantage estimation: agents within the same functional group share a group-level mean-return baseline, enabling lower-variance updates and stronger intra-group coordination. Across representative attack scenarios, AdvEvo-MARL consistently keeps attack-success rate (ASR) below 20\%, whereas baselines reach up to 38.33\%, while preserving or even improving task accuracy (up to +3.67\%). These results show that safety and utility can be jointly improved without relying on extra guard agents or added system overhead.

Poster
Learn to change the world: Multi-level reinforcement learning with model-changing actions

Ziqing Lu, Babak Hassibi, Lifeng Lai, Weiyu Xu

Expand abstract

Reinforcement learning usually assumes a given or sometimes even fixed environment in which an agent seeks an optimal policy to maximize its long-term discounted reward. In contrast, we consider agents that are not limited to passive adaptations: they instead have model-changing actions that actively modify the RL model of world dynamics itself. Reconfiguring the underlying transition processes can potentially increase the agents' rewards. Motivated by this setting, we introduce the multi-layer configurable time-varying Markov decision process (MCTVMDP). In an MCTVMDP, the lower-level MDP has a non-stationary transition function that is configurable through upper-level model-changing actions. The agent's objective consists of two parts: Optimize the configuration policies in the upper-level MDP and optimize the primitive action policies in the lower-level MDP to jointly improve its expected long-term reward.

Poster
Trojan-Speak: Bypassing Constitutional Classifiers with No Jailbreak Tax via Adversarial Finetuning

Bilgehan Sel, Xuanli He, Alwin Peng, Ming Jin, Jerry Wei

Expand abstract

Fine-tuning APIs offered by major AI providers create new attack surfaces where adversaries can bypass safety measures through targeted fine-tuning. We introduce **Trojan-Speak**, an adversarial fine-tuning method that bypasses Anthropic's Constitutional Classifiers. Our approach uses curriculum learning combined with GRPO-based hybrid reinforcement learning to teach models a communication protocol that evades LLM-based content classification. Crucially, while prior adversarial fine-tuning approaches report more than 25\% capability degradation on reasoning benchmarks, Trojan-Speak incurs less than 5\% degradation while achieving 99+\% classifier evasion for models with 14B+ parameters. We demonstrate that fine-tuned models can provide detailed responses to expert-level CBRN (Chemical, Biological, Radiological, and Nuclear) queries from Anthropic's Constitutional Classifiers bug-bounty program. Our findings reveal that LLM-based content classifiers alone are insufficient for preventing dangerous information disclosure when adversaries have fine-tuning access, and we show that activation-level probes can substantially improve robustness to such attacks.

Spotlight
Structured Expert Routing with Multi-View Task Priors for Offline Meta-Reinforcement Learning

Yisen Zhao, Peixi Peng, Xinyu Hu, Cong Li, Zhan Su, Zhuojian Li

Expand abstract

Offline meta-reinforcement learning requires agents to generalize to unseen tasks from fixed datasets, yet existing sequence-based and MoE-based methods rely on implicit or token-level routing signals that fail to capture task-level structure. We propose the **Task-Guided Router (TGR)**, a structured expert-routing framework that explicitly models inter-task relationships via multi-view task representations that combine semantic descriptors, behavioral summaries, and latent dynamics features. Using structure-guided routing, TGR assigns experts based on global task compatibility rather than local trajectory fragments, enabling stable specialization and effective knowledge transfer across tasks.Extensive experiments on continuous-control benchmarks demonstrate that TGR consistently outperforms state-of-the-art offline meta-RL methods in few-shot generalization, particularly under sparse data and heterogeneous dynamics. Our results highlight the importance of task-level priors for robust offline meta-reinforcement learning.

Poster
Contrastive Reasoning Alignment: Reinforcement Learning from Hidden Representations

Haozheng Luo, Yimin Wang, Jiahao Yu, Binghui Wang, Yan Chen

Expand abstract

We propose CRAFT, a red-teaming alignment framework that leverages model reasoning capabilities and hidden representations to improve robustness against jailbreak attacks. Unlike prior defenses that operate primarily at the output level, CRAFT aligns large reasoning models to generate safety-aware reasoning traces by explicitly optimizing objectives defined over the hidden state space. Methodologically, CRAFT integrates contrastive representation learning with reinforcement learning to separate safe and unsafe reasoning trajectories, yielding a latent-space geometry that supports robust, reasoning-level safety alignment. Theoretically, we show that incorporating latent–textual consistency into GRPO eliminates superficially aligned policies by ruling them out as local optima. Empirically, we evaluate CRAFT on multiple safety benchmarks using two strong reasoning models, Qwen3-4B-Thinking and R1-Distill-Llama-8B, where it consistently outperforms state-of-the-art defenses such as IPO and SafeKey. Notably, CRAFT delivers an average **79.0%** improvement in reasoning safety and **87.7%** improvement in final-response safety over the base models, demonstrating the effectiveness of hidden-space reasoning alignment.

Poster