--- title: 'CogER-Agent: Elastic Reasoning for LLMs' url: https://www.emergentmind.com/topics/coger-agent type: topic --- # CogER-Agent: Elastic Reasoning for LLMs CogER-Agent is a cognitive-inspired reinforcement learning agent designed for elastic reasoning in large language models (LLMs). Developed within the CogER (Cognitive-Inspired Elastic Reasoning) framework, the agent autonomously selects among multiple reasoning strategies per query, adapting computational effort according to task complexity. CogER-Agent models hierarchical human reasoning as a Markov Decision Process (MDP) and is trained with a Group Relative Policy Optimization (GRPO) algorithm. It supports tool-assisted reasoning, allowing the LLM to invoke external resources as necessary, and demonstrates significant efficiency and accuracy gains over prior scaling and routing methods [2512.15089]. ## 1. Markov Decision Process Formulation CogER-Agent's reasoning strategy selection is formalized as a finite-horizon MDP, $(\mathcal{S}, \mathcal{A}, \mathcal{T}, \mathcal{R}, \pi)$. The state at time $t$ is $s_t = [x, y_{1:t-1}, L_i]$, where $x$ is the user query, $y_{1:t-1}$ the partial chain-of-thought (CoT), and $L_i$ the inferred complexity level ($L_i \in \{L_1, L_2, L_3, L_4\}$). Each state either contains an explicit complexity tag or maintains the information within the model's hidden state. The action space $\mathcal{A}$ comprises four discrete high-level modes: - $\mathrm{No\_Think}$: Immediate answer (level-1, rapid inference) - $\mathrm{Think}$: Intermediate reasoning (e.g., invoking a larger LLM) - $\mathrm{Extend}$: Escalated chain-of-thought - $\mathrm{Delegate}$: External tool invocation and the vocabulary $\mathcal{V}$ for standard token emission, so $\mathcal{A} = \{\mathrm{No\_Think}, \mathrm{Think}, \mathrm{Extend}, \mathrm{Delegate}\} \cup \mathcal{V}$. The transition function $\mathcal{T}(s_t, a_t)$ deterministically updates the state, either by appending a mode-selection tag or emitting a token. The reward function is a terminal scalar combining format validation, output correctness, and penalization for unnecessary cognitive escalation: $$ \mathcal{R}(s, a) = \mathcal{R}_{\mathrm{format}}(s, a) + \mathcal{R}_{\mathrm{accuracy}}(s, a) + \mathcal{R}_{\mathrm{hierarchy}}(s, a). $$ $\mathcal{R}_{\mathrm{format}}$ enforces well-formed tags, $\mathcal{R}_{\mathrm{accuracy}}$ rewards correct answers, and $\mathcal{R}_{\mathrm{hierarchy}}$ is defined as $b(L_{\min}(s)) - \delta(L_{\min}(s), L(s))$ with $b(L)=0.5\,(L-1)$ and $\delta(L_{\min},L) = 0.2\,(L - L_{\min})_+$ to encourage minimal sufficient reasoning. ## 2. Policy Parameterization and Group Relative Policy Optimization The policy $\pi_\theta(a_t \mid s_t)$ is a single-headed softmax over joint actions: $$ \pi_\theta(a_t \mid s_t) = \frac{\exp(f_\theta(s_t)_a)}{\sum_{a'} \exp(f_\theta(s_t)_{a'})} $$ where $f_\theta$ outputs unnormalized logits. CogER-Agent is optimized using Group Relative Policy Optimization (GRPO), a PPO-style algorithm with KL regularization. For each group of $G$ rollouts $\{o_i\}_{i=1}^G$ under the old policy $\pi_{\theta_{\mathrm{old}}}$, rewards $r_i$ are computed and normalized as $\tilde r_i = (r_i - \mathrm{mean}(r)) / \mathrm{std}(r)$, giving relative advantage $\hat A_i = \tilde r_i$ to each token in $o_i$. The GRPO objective is: $$ \mathcal{J}_{\mathrm{GRPO}}(\theta) = \mathbb{E}_{x, \{o_i\}} \left[ \frac{1}{G} \sum_{i=1}^G \min\left\{ \frac{\pi_\theta(o_i \mid x)}{\pi_{\theta_{\mathrm{old}}}(o_i \mid x)} \hat{A}_i, \mathrm{clip}\left(\frac{\pi_\theta(o_i \mid x)}{\pi_{\theta_{\mathrm{old}}}(o_i \mid x)}, 1-\varepsilon, 1+\varepsilon\right)\hat{A}_i \right\} - \beta D_{\mathrm{KL}}[\pi_\theta \Vert \pi_{\mathrm{ref}}] \right] $$ where $\varepsilon$ is the PPO clipping threshold, $\beta$ the KL penalty, and $\pi_{\mathrm{ref}}$ a small reference policy. ## 3. Training Procedure Query difficulty is estimated in situ by the 7B-parameter CogER-Agent using a prompt requiring level assignment ($L_1$–$L_4$) prior to answer generation. For $L_1$, a direct answer is also output. For each query, $G=12$ complete trajectories (rollouts) are sampled, collectively labeled with the terminal reward as in equations (1)–(4) and normalized within-group to yield the advantage. Fine-tuning employs GRPO combined with LoRA adaptation (rank $r=16$), using AdamW ($5 \times 10^{-5}$ learning rate), batch size $24 \times 3$, and a maximum generation of $8192$ tokens. Training is performed over a single epoch using an 8,000-query mixture (2,000 each from GSM8K, MATH, CommonsenseQA, MedQA), with three random seeds ($\{21,26,42\}$) and mean ± standard deviation reported over three runs. ## 4. Cognitive Tool-Assisted Reasoning (CoTool) For level-4 (Delegate) queries, CogER-Agent invokes a tool-augmented reasoning mechanism ("CoTool"). The agent alternates between internal generation and external tool invocation: - Generation proceeds until the LLM emits EOS or a tool query boundary token. - On emission of `<|end_tool_query|>`, the sub-query $q_{\mathrm{tool}}$ is extracted. - A dedicated tool selector, prompted with task-specific instructions, chooses the optimal RSTKit tool and JSON arguments, executes the tool, and returns the output. - The tool's result is appended to the CoT (as `<|begin_tool_result|>...`), resuming internal reasoning. - The process is iterated up to a maximum number of tool calls/turns. Formally, the probability of an $i$-th tool query is given by: $$ P\left(q_{\mathrm{tool}}^{(i)} \mid I, q, R^{(i-1)}\right) = \prod_{t=1}^{T_q^{(i)}} P\left(q_{\mathrm{tool},t}^{(i)} \mid q_{\mathrm{tool},