Training Language Models to Cooperate with Inference-Time Controllers

Moumita Choudhury Vanshaj Khattar Jing Liu T. Koike-Akino Ankush Chakrabarty S. Zilberstein Ye Wang

Jul 2026 · arXiv.org · Vol abs/2607.23771 · 0 citations · 34 references

Computer Science

TL;DR

CALM (Controller-Aware Language Models), a post-training framework that explicitly places controllers in the training loop, is introduced, showing that controller-aware post-training improves generalization across inference-time workflows beyond single-controller optimization.

Abstract

Large language model (LLM) performance increasingly depends not only on the base model, but also on the inference-time controller used to organize reasoning. Existing post-training methods, however, typically optimize for a single fixed interaction pattern, despite real deployments relying on diverse controllers such as Chain-of-Thought, self-consistency, debate, planning, and verification pipelines. This creates a training--deployment mismatch and limits transfer to new workflows. We introduce CALM (Controller-Aware Language Models), a post-training framework that explicitly places controllers in the training loop. We formulate controller-aware post-training as multi-task reinforcement learning over controller-induced interaction protocols, where controllers are compositions of reusable local reasoning modules. This structure also induces a module-level decomposition of mixed-controller training under a turn-level GRPO objective, enabling a systematic study of controller and module-aware training strategies. We evaluate CALM on held-out controller compositions and broader controller shifts, showing that controller-aware post-training improves generalization across inference-time workflows beyond single-controller optimization.

DOI arXiv Semantic Scholar

View source

Similar papers

Preprint Aug 2026

Performance Foundations of Parallel&Distributed Reasoning Language Models

This work systematize the RL-for-LLM paradigm and provides a compute-centric analysis of prominent post-training algorithmic frameworks: Proximal Policy Optimization (PPO), Group Relative Policy Optimization (GRPO), as well as their variants, and develops a taxonomy of intra- and inter-model parallelism strategies for...

Maciej Besta, L. Schmidt, Lara Nonino et al. · 0 citations

Preprint Aug 2026

SPT: Skills as Pre-Training Data for Agentic Language Models

Agentic (tool-using) language models are mainly trained on tool-call traces and agent trajectories during post-training. These data provide direct behavioral supervision, but producing them requires task environments, execution, and verification, making broad tool and task coverage expensive. Publicly available skills...

Yufei Sun, Yudong Li, Yi-Min Cheng · 0 citations

Jul 2026

TAPO: Transition-Aware Policy Optimization for LLM Agents

This work proposes TAPO: Transition-Aware Policy Optimization for LLM Agents, a unified training framework that alternates between policy optimization and transition supervision, and demonstrates that TAPO consistently improves task performance over pure policy optimization baselines.

Cong Li, Pei-Xi Peng, Yisen Zhao et al. · arXiv.org · 1 citation

Preprint Aug 2026

MCP-Universe RL: A Framework for Training MCP Tool-Use Agents via Reinforcement Learning

Reinforcement learning (RL) has become an effective way to improve the tool-use ability of large language models (LLMs), but most existing RL frameworks stop at the policy update. For every new domain, the user is left with two hard systems problems: standing up an isolated environment for each of hundreds of concurren...

Ziyang Luo, Yan Yang, Xiang-Ru Jian et al. · 0 citations

Preprint Aug 2026

Cognitive Demand Steering for Adaptive Meta-Reasoning in Large Language Models

CDS is introduced, a training-free meta-reasoning framework equipped with residual demand assessment: at each step, an LLM-based progress evaluator characterizes the residual reasoning required to arrive at a solution rather than merely evaluating the previous step.

John Scoville, Shengzhuang Chen, Yejin Bang et al. · 0 citations

Book Open access Aug 2026

Large Language Model (LLM) as an Excellent Reinforcement Learning Researcher in both Single-Agent and Multi-Agent Scenarios

A Self-Evolutional single-agent/multi-agent Reinforcement Learning (SE-RL) framework that utilizes a Large Language Model (LLM) to design various RL algorithm modules, such as agent model design, reward function, profiling, communication, and state imagination, by leveraging the LLM generating module output or code.

Vincent Fu, Xin-Xin Xu, Weichen Xu et al. · Proceedings of the 32nd ACM... · 0 citations