Paper deep dive

A Graph-Based Reinforcement Learning Framework for Structured Drift Diagnosis and Recovery in Autonomous LLM Agents

Ismail El Hamraoui, Sagar Jose, Nicolas Bureau, Robert Plana

Year: 2026Venue: arXiv preprintArea: cs.AIType: PreprintEmbeddings: 61

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 94%

Last extracted: 8/17/2026, 5:05:18 AM

Summary

This paper introduces a graph-based reinforcement learning framework for diagnosing and recovering from behavioral drift in autonomous LLM agents. The system uses a small language model specialized via Group Relative Policy Optimization (GRPO) to traverse a five-node recovery graph. Each node performs a specific diagnostic role (drift classification, operation detection, documentation search, risk evaluation, and final decision) using structured XML reasoning. The framework aims to detect when an agent deviates from its task, assess the reversibility of actions taken during the drift, and execute appropriate recovery strategies such as message history rollback or human escalation, without retraining the main large language model.

Entities (17)

Behavioral Drift · problem · 99%Ismail El Hamraoui · person · 99%Sagar Jose · person · 99%Graph-Based Reinforcement Learning Framework · methodology · 98%GRPO · algorithm · 97%Small Language Model · model · 96%AppWorld · benchmark · 95%Assystem EOS · organization · 95%Recovery Graph · architecture · 94%Structured XML Reasoning · output-format · 93%Drift Classification · task · 92%Operation Detection · task · 92%Risk Evaluation · task · 92%LLM-as-judge · technique · 90%Type I Drift · drift-category · 88%Type II Drift · drift-category · 88%Type III Drift · drift-category · 88%

Relation Signals (15)

Ismail El Hamraoui → affiliatedwith → Assystem EOS

confidence 99% · Ismail El Hamraoui*Sagar Jose†Nicolas Bureau Robert Plana Digital Excellence Center, Assystem EOS, France

Sagar Jose → affiliatedwith → Assystem EOS

confidence 99% · Ismail El Hamraoui*Sagar Jose†Nicolas Bureau Robert Plana Digital Excellence Center, Assystem EOS, France

Graph-Based Reinforcement Learning Framework → addresses → Behavioral Drift

confidence 98% · Autonomous LLM agents are increasingly deployed in complex real-world workflows, yet they remain vulnerable to runtime behavioral drift... this work targets a plug-and-play recovery module instead.

Graph-Based Reinforcement Learning Framework → employs → Small Language Model

confidence 96% · It introduces a graph-based framework in which a single small language model is trained via reinforcement learning

Graph-Based Reinforcement Learning Framework → evaluatedon → AppWorld

confidence 95% · Experiments on the public AppWorld benchmark show that the method generally exploits information about the suspected drift onset

Graph-Based Reinforcement Learning Framework → uses → GRPO

confidence 95% · The same model can be trained on all five roles at once with GRPO.

Graph-Based Reinforcement Learning Framework → implements → Recovery Graph

confidence 94% · It introduces a graph-based framework in which a single small language model is trained via reinforcement learning to specialize at each node of a recovery graph

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:Autonomous LLM agents are increasingly deployed in complex real-world workflows, yet they remain vulnerable to runtime behavioral drift, a silent deviation from the original task that can lead to irreversible side effects on external systems. Existing approaches address drift at the prompt level but lack structured mechanisms for step-level detection, risk assessment, and recovery decision. Because the main task-executing agent is often a large and expensive model that cannot be re-trained on every deployment, this work targets a plug-and-play recovery module instead. It introduces a graph-based framework in which a single small language model is trained via reinforcement learning to specialize at each node of a recovery graph, external to the main agent. Each node has a precise role\,: drift classification, operation detection, risk evaluation, or final decision and the model learns to produce structured XML-formatted reasoning adapted to that role. Training combines rule-based structural rewards with an LLM-as-judge semantic-quality signal, so that the model is graded both on how it answers (schema and length) and on what it says. Experiments on the public AppWorld benchmark show that the method generally exploits information about the suspected drift onset to issue correct recovery decisions using a small language model. In addition, the trained small language model reliably respects the prescribed output schema and produces semantically appropriate content in each field according to its assigned node role.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

57,732 characters extracted from source content.

Expand or collapse full text

A Graph-Based Reinforcement Learning Framework for Structured Drift Diagnosis and Recovery in Autonomous LLM Agents Ismail El Hamraoui*Sagar Jose†Nicolas Bureau Robert Plana Digital Excellence Center, Assystem EOS, France Abstract Autonomous LLM agents are increasingly deployed in complex real-world workflows, yet they remain vulnera- ble to runtime behavioral drift, a silent deviation from the original task that can lead to irreversible side ef- fects on external systems. Existing approaches address drift at the prompt level but lack structured mecha- nisms for step-level detection, risk assessment, and re- covery decision. Because the main task-executing agent is often a large and expensive model that cannot be re-trained on every deployment, this work targets a plug-and-play recovery module instead. It introduces a graph-based framework in which a single small lan- guage model is trained via reinforcement learning to specialize at each node of a recovery graph, external to the main agent. Each node has a precise role : drift classification, operation detection, risk evaluation, or final decision and the model learns to produce struc- tured XML-formatted reasoning adapted to that role. Training combines rule-based structural rewards with an LLM-as-judge semantic-quality signal, so that the model is graded both on how it answers (schema and length) and on what it says. Experiments on the public AppWorld benchmark show that the method generally exploits information about the suspected drift onset to issue correct recovery decisions using a small lan- guage model. In addition, the trained small language model reliably respects the prescribed output schema and produces semantically appropriate content in each field according to its assigned node role. Keywords: LLM Agents, Behavioral Drift, Reinforce- ment Learning, Graph-based Recovery, Small Language Models, Structured Reasoning. Manuscript submitted to Applied Intelligence (Springer) for peer review. Preprint version.†Corresponding author: sajose@assystem.com. *First author. 1 Introduction LLM agents no longer just answer questions but they act. They call APIs, send emails, transfer money, and modify databases on behalf of their user. Frameworks like ReAct [Yao et al., 2023] popularized this pattern, and benchmarks such as AppWorld [Trivedi et al., 2024] now test agents through long trajectories of tool calls against hundreds of APIs. As agents do more real actions on real systems, each mistake costs more. One kind of mistake stands out: behavioral drift. Drift is not a crash nor merely a hallucinated answer. Rather, a hallucinated answer may trigger a sequence of actions that ultimately leads to behavioral drift. It is the mo- ment when the agent begins to lose sight of the initial goal. This can happen because a malicious instruction was hidden in a webpage or an email that the agent happened to read [Greshake et al., 2023], or simply because the agent over-extended a vague sub-goal on its own. Once drift starts, The agent may continue executing subsequent steps with confidence, while pro- gressively drifting away from the initial goal. And since agents now perform write operations (sending money, deleting messages, changing settings), a drifted agent can leave the environment in a state that resetting the model’s memory will not fix: rolling back the message history does not roll back the world [Zheng, 2025]. In some cases, the agent cannot complete the original task without either executing an explicit correction action or escalating to a human operator so that the latter can repair the environment before execution resumes. Most existing work on drift focuses on preventing it at the input, for example by isolating untrusted content or by training an instruction hierarchy [Wallace et al., 2024]. The post-drift action is still missing. A procedure is needed that (i) finds when the drift started by walking the trajectory backwards, (i) lists the write operations the agent actually committed during drift, (i) decides 1 arXiv:2608.14109v1 [cs.AI] 14 Aug 2026 for each of them whether it can be undone with an inverse API call, and (iv) chooses between rolling back the message history and escalating to a human. Asking a single LLM to do all four inside one prompt is brittle, and free-form chain-of-thought [Wei et al., 2022] does not give the schema discipline that each of these sub- tasks needs. This paper introduces a graph-based framework that turns drift recovery into a small state machine of five nodes, where a single small LLM is specialized at each node via reinforcement learning. Every node emits the same simple format, one<reasoning>block followed by one<answer>JSON block and the same model can be trained on all five roles at once with GRPO. The main contributions of this work are: • A node-level fine-tuning methodology, in which a single small LLM is specialized across multiple roles through role-conditioned prompts and a shared GRPO objective, without training a separate policy per node. • A graph-based approach to drift recovery in LLM agents, framing recovery as a routed traversal of a diagnostic state machine that walks the trajectory backwards until it finds an aligned streak of steps. 2 Related Work 2.1Attacks and Defenses on LLM Agents A lot of recent work looks at how tool-using LLM agents can be pushed away from their task. The main mecha- nism is prompt injection: attacker-controlled text hidden inside a webpage, an email, or a tool response takes over the agent [Greshake et al., 2023]. Benchmarks such as AgentDojo [Debenedetti et al., 2024] now measure how often this succeeds, and how often it also breaks the original task. On the defense side, most proposals stop the injection before it reaches the model: isolating untrusted content in a separate LLM [Willison, 2023], teaching the model an instruction hierarchy [Wallace et al., 2024], or placing a firewall between the agent and its tools [Debenedetti et al., 2025]. All of these target prevention. The approach assumes drift has already happened, because in practice some fraction of injections and vague sub-goals will slip through. As Zheng [2025] point out, resetting the model’s message history does not undo any environment write the agent has already made. So the problem this work tackles is: given that drift happened, walk the trajectory back, look at what was actually written, and decide what to do. 2.2 Graph-based Agents Most LLM agents run in a single reasoning-acting loop like ReAct [Yao et al., 2023]. Frameworks like Lang- Graph [LangChain Team, 2024] extend this pattern to explicit multi-node state machines, and this work builds on that. Two things make the graph proposed here different. First, each of its five nodes plays a specific role (classify, extract, look up, evaluate, decide), instead of repeating the same Thought/Action/Observation loop. Second, the graph is diagnostic rather than task-solving: it does not push a task forward, it goes back in time to reconstruct what happened during a drifted region of the trajectory. 2.3 Reinforcement learning for LLMs Training is carried out with GRPO [Shao et al., 2024], a variant of PPO that drops the value network and computes advantages relative to a group of rollouts of the same prompt. DeepSeek-R1 [Guo et al., 2025] showed that with rule-based verifiable rewards, this is enough to elicit long reasoning without any human preference data. GRPO fits the setting of this work for three simple reasons: (i) the reward has verifiable parts (Does the output parse without error, and do the JSON keys match?), (i) the group-relative signal stays useful even when all completions are bad early in training, and (i) no value network means training can be performed on a single GPU. 2.4 Structured Reasoning Chain-of-Thought [Wei et al., 2022] showed that letting a model write intermediate steps before answering helps. Reflexion [Shinn et al., 2023] pushes this further by having the agent write a natural-language critique of its own past attempt and use it as a hint for the next try. The<reasoning>+<answer>envelope used in this work is a deliberate move away from free-form CoT. In the proposed pipeline, the output of one node is the input of the next. If each node writes free-form text, the next one has to reparse it and hope to extract the same structure. Forcing a fixed schema turns each node’s output into a typed message that the next node can just read. 2 3 Problem Formulation 3.1 Notation and Setup LetTdenote a task executed by an autonomous LLM agent. The execution of T produces a trajectory τ = (s 0 ,s 1 ,...,s n ),(1) where each states i captures the agent’s internal state (context, plan, next-action decision) together with the interaction actually performed at stepi. The agent interacts with a set of external environment entities E = e 1 ,e 2 ,...,e m ,(2) where an entity may correspond to an API, a database, a file system, a software application, or any other external service exposed as a tool call to the agent. Each step’s interaction withEdecomposes into a set of read operationsR i (observations ofEthat leave it unchanged) and a set of write operationsW i (operations that mutate E). We denote the per-step interaction by x i = (R i ,W i ) and the full environment-facing execution trace by X (τ ) = x 0 ,x 1 ,...,x n .(3) 3.2 Ideal Executions and Task Success For a given taskT, there is in general no unique task- faithful trajectory. Two executions may satisfyTwhile differing in the order of independent read operations, in the decomposition of a sub-goal into a longer or shorter sequence of tool calls, or in the specific API path chosen to commit the same final environment state. We therefore model the ideal execution not as a single trace but as a set of task-faithful trajectories, G(T ) = X (1) ,X (2) ,... ,(4) each elementX (j) =g (j) 0 ,g (j) 1 ,... being a valid ground-truth trace withg (j) i = (R ⋆,(j) i ,W ⋆,(j) i ). Corre- spondingly, the task is considered successfully completed if there exists at least one X (j) ∈G(T ) such that X (τ ) ≡ X (j) ,(5) where≡denotes semantic equivalence up to reorder- ing of independent reads and any bijective substitution of API paths that yields the same final environment state. Eq. 5 is what benchmarks such as AppWorld’s task_goal_completion[Trivedi et al., 2024] actually measure in practice; the definitions that follow do not depend on the specific choice of equivalence relation. Two consequences follow. First, matching a per-step idealg ⋆ i is not the right primitive: a stepx i that differs fromg (j) i but coincides withg (j ′ ) i for some otherj ′ is still task-faithful. Second, the notion of drift must be relative to someX (j) ∈G(T), not to a single canonical trace. 3.3 Behavioral Drift Definition 1 (Behavioral Drift). Given the setG(T), we say that drift occurs at stepdif there exists a task- faithful referenceX (j) ∈G(T) such that the execution matches it before d, x i ≡ g (j) i ∀i < d,(6) and, from stepdonward, no such reference can be maintained: ∀X (j ′ ) ∈G(T ), ∃i≥ d, x i ̸≡ g (j ′ ) i .(7) The step d is the drift onset. Remark 1 (Drift is not monotone after d). Eq. 7 does not imply that every step afterdis off-task.There are cases where drifted regions are interleaved: after the onset, the agent may temporarily return to task- consistent behavior, for instance by accidentally query- ing a correct API while pursuing an attacker’s sub-goal that happens to overlap with a legitimate one, before drifting again a few steps later. A realistic post-dpat- tern may therefore look like D→ D→ D | z drifted → A→ A→ A |z aligned → D→· rather than a clean tail of off-task steps. Given a suspected onset ˆ d, the goal of recovery is then to produce a recovered trajectory τ r such that X (τ r ) ≡ X (j) ∈G(T ),(8) i.e., the final execution is equivalent to some task-faithful reference. Whether Eq. 8 is achievable at all, and if so at what cost, depends not only on when drift started, but crucially on what the agent did between d and n. Definition 2 (Irreversible Operation). A write w ∈ W i is irreversible if no operation is available to the agent at recovery time that restores the state ofE observed immediately beforewusing only information available to the agent at that moment. 3.4 Drift Taxonomy Under the working definitions above, we distinguish three drift categories based on the extent of their impact on the agent’s operating environment while the agent is in drift region [d,n] and on the causal persistence of the drift source. The categories are mutually exclusive and determine which recovery strategy is admissible. 3 Type I - Transient Read Drift. The drifted region produces no environment mutation, W i =∅ ∀i≥ d,(9) and its causal source is transient: a one-off stochastic misstep, a spurious hallucination not anchored in a persistent input. Re-executing from a checkpoint prior todis therefore unlikely to reproduce the same failure, and a plain message-history rollback is sufficient to obtain τ r . Type I - Persistent Read Drift. The drifted region still involves only read operations, W i =∅ ∀i≥ d,(10) but its causal source persists in the environment or in the task inputs: a poisoned document still sitting in the mailbox, an untrusted webpage that will be fetched again, a compromised knowledge base still returning the same injected instruction. Re-executing from before dtherefore re-exposes the agent to the same trigger, and naive rollback is insufficient. Recovery additionally requires that the agent retain knowledge of the encoun- tered drift through a dedicated recovery memory or an augmented context passed to the re-executing agent, enabling future decisions to be informed by previous failures, consistent with self-reflective paradigms such as Reflexion [Shinn et al., 2023]. Type I - Write Drift. During the drifted region, the agent performs at least one write, ∃i≥ d, W i ̸=∅,(11) so the environment has been mutated: E ′ ̸=E.(12) Rollback of the agent state alone cannot re-establish the equivalenceX(τ r )≡ X (j) required by Eq. 8: the message history resets, but the world does not [Zheng, 2025]. Recovery from Type I drift jointly requires (i) repairing or compensating each environment modifica- tionwthat admits a documented inverse operation, and (i) rolling back the agent state and re-executing the remaining trajectory from a pre-drift checkpoint. When even a single writew ∈ S i≥d W i is irreversible in the sense of Definition 2, no purely automated recovery is safe, and the incident must be escalated to a human operator. 3.5 Scope of This Work This paper focuses primarily on Type I drift, where the environment remains unchanged and message-history rollback is a correct recovery. For Type I drift, the same rollback mechanism is used, together with a drift- aware recovery memory that enriches the re-executed context so that the initial agent does not fall for the same drift source a second time. For Type I drift, the framework does not attempt to compensate environ- ment modifications automatically; instead, it focuses on detecting whether the drifted writes are reversible or ir- reversible, so that a rollback is authorized only when the drifted region contains read operations exclusively, and the incident is escalated to a human operator as soon as any write has been committed. In future work, this dichotomy will be handled asymmetrically: reversible writes will be corrected in place by the recovery agent through the corresponding inverse API calls, after which task execution resumes, and only genuinely irreversible writes will trigger human escalation. Executing these in- verse API calls, and closing the loop into an automated correction step, is deferred to future work. Finally, drift is assumed monotone once started: ev- ery step after the onsetdis off-task, ruling out the interleaved pattern of Remark 1. 3.6 Recovery Acceptance Criteria Before turning to the proposed method, this section lists the properties that guided the design of the framework introduced in this work. Fig. 1 summarizes them. •Drift-robust. The recovery module must never be biased or influenced by the drift itself. If the same malicious content that pushed the main agent off-task can also push the recovery module off-diagnosis, the whole procedure collapses. •Diagnostic Autonomy. It must identify the drift’s root cause and the exact step at which drift started, using only information already available at recovery time. No oracle, no ground-truth trace. •State-aware. Rollback alone is not a robust solu- tion: resetting the message history does not reset the environment [Zheng, 2025]. Any real recovery must reason about the environment state, not just the agent’s context. •Action reversibility. It must correct reversible ac- tions when possible, and escalate irreversible actions to a human with enough context to act. A single uni- form response (always rollback, or always escalate) is either unsafe or uselessly conservative (Out of scope). •Classification reversibility. It must be able to dis- tinguish between reversible and irreversible actions, because that distinction is what makes the dual strat- egy above decidable. 4 Figure 1: Desiderata for a drift-recovery framework in a deployed LLM-agent setting. •Plug-and-play. It must integrate into an existing agent without retraining the main task-executing model, which is typically large, expensive, and shared across many deployments. •Resource-efficient. It must remain low-cost in terms of compute and memory. •Preventive. It should reduce the probability of drift att+Twhen the pattern seen attrecurs, e.g. by surfacing the incident structure back to the outer system for future guardrails. 4 Methodology 4.1 Overview Our approach models drift recovery as a routed traversal of a small state machine. All five nodes call the same policy but with role-specialized prompts, so the model plays a different role at each stop of the graph. 4.2 Recovery Graph Architecture The recovery graph, shown on the left of Fig. 2, consists of five nodes: 1.Classify Drift (n 1 ): given the task, a suspected drift onset supplied by an external drift detector, and the content of a single step, decides whether the step is aligned with the task (Y) or drifted (N), and gives a short reason. Note thatn 1 does not perform autonomous drift detection: it refines an externally provided detection signal into a per-step verdict. 2.Detect Operations (n 2 ): triggered only whenn 1 flags a step as drifted. Enumerates every successfully executed write operation and every off-task read per- formed in that step. Failed writes and task-aligned reads are excluded. 3. Search Documentation (n 3 ): given the list of writes, outputs the deduplicated set of applications involved (e.g.spotify | gmail) so their documen- tation can be fetched. 4.Evaluate Risk (n 4 ):given the writes and the fetched API documentation, parti- tions the writes intowrite_reversibleand write_not_reversible. 5.Aggregate & Decide (n 5 ): given the full per-step diagnosis, produces the final action rollback_before_driftwith the target step, or escalate_human with a structured incident report. 5 Figure 2: Training pipeline. The five-node recovery graph (left) emits a node-tagged promptqat each audited step. The policyπ θ samplesGcompletions per prompt in the shared⟨reasoning⟩+⟨answer⟩XML schema , each scored by the composite rewardR(o,ν,q) = P 9 k=1 w k r k (o,ν,q) that combines eight structural checks with a frozen LLM-as-judge signal . Group-normalized advantages ˆ A i drive the GRPO update, specializingπ θ into all five node roles. 4.3 Structured XML Reasoning Every node forces the model to emit exactly one <reasoning>block followed by exactly one<answer> block, the latter containing valid JSON with a node- specific key set: 1 <reasoning > 2 analysis 3 </reasoning > 4 5 <answer > 6 "<node -specific -keys >": ... 7 </answer > Listing 1: Shared output schema; keys inside<answer> are node-specific. The expected key set per node is: • n 1 : step, is_aligned, why • n 2 : write_operations, read_operations • n 3 : apps_name • n 4 : write_reversible,write_not_reversible • n 5 : action, arguments 4.4Backward Walk and Aligned-Streak Termination The routing rule between the five nodes implements a simple but consequential control law. After n 1 : • If is_aligned = N, route to n 2 . • Else, if the lastKalignment verdicts are allY, route to n 5 (recovery termination). •Else, decrement the pending step index by 1 and loop back to n 1 (backward walk). Aftern 2 : ifwrite_operations = None(only off-task reads, or nothing successful), route back ton 1 ; otherwise route ton 3 . Aftern 3 : unconditional edge ton 4 . After n 4 : unconditional edge back to n 1 . The aligned-streak lengthKreflects the observa- tion that a single aligned step immediately preceding a drifted region is not sufficient evidence that drift has ended; a short streak is a robust and cheap proxy. 6 Algorithm 1 The approach Require:taskT, trajectoryH, suspected onsetk 0 , streak length K 1: k ← k 0 ; w ← [ ]; steps← [ ] 2: while k ≥ 1 do 3: s k ← GetStep(k) 4:(aln, why)← n 1 (T ,k 0 ,s k ) 5:append aln to w; append diagnosis to steps 6: if aln = N then 7:(O w ,O r )← n 2 (s k ) 8:if O w ̸=∅ then 9:P ← n 3 (O w ) 10:D ← FetchDocs(P) 11:(O rev w ,O irr w )← n 4 (O w ,D) 12:end if 13: else if |w|≥ K and last K of w are all Y then 14:break 15: end if 16: k ← k− 1 17: end while 18: d, arg← n 5 (T ,k 0 , steps) 19: if d = rollback_before_drift then 20:RollbackMessages(arg.rollback_to_step) 21: else 22:EscalateHuman(arg.report) 23: end if 4.5Node-Specialized RL Training with GRPO The five roles introduced in the recovery graph share the same modelπ θ ; what changes across nodes is the prompt schema and the shape of the expected<answer>JSON. This work adopts Group Relative Policy Optimization (GRPO) [Shao et al., 2024, Guo et al., 2025] to specialize π θ at every node with a shared rollout procedure. Group rollouts on structured prompts. For every training promptq(drawn from a node-tagged pool of in- spection transcripts, with tagν ∈n 1 ,n 2 ,n 3 ,n 4 ,n 5 ), a group ofGcompletionso (1) ,...,o (G) is sampled from an old policyπ θ old . Eacho (i) is a full XML response, i.e. a<reasoning>block followed by an<answer>JSON block. The XML schema plays a central role here: be- causeqends with the strict format instructions of the corresponding node prompt, well-behaved samples fall inside a narrow parseable region, while ill-behaved sam- ples land outside it. GRPO exploits this contrast di- rectly, without training a value network, by comparing samples inside a group. Composite reward. Each completion receives a scalar rewardr i =R(o (i) ,ν,q) from a composite func- tion that combines rule-based structural checks with a semantic-quality signal produced by a frozen LLM-as- judge (Section 4.6). No reward model is trained; the judge is used at inference time only. This design extends the R1-style verifiable-reward recipe [Guo et al., 2025] with an additional non-parametric grader that closes the gap between form and content. Group-normalized advantages. Advantages are computed relative to the group only: ˆ A i = r i − ̄r std(r) + ε ,(13) with ̄randstd(r) the group mean and standard devia- tion. This removes the need for a value function and yields a stable signal even whenr i concentrates on a few components of the composite reward. Objective. Denoting byπ θ (o| q) the token-level pol- icy and byπ ref a frozen reference (the pretrained base), the GRPO objective is J (θ) =E q,o i h min ρ i ˆ A i , clip(ρ i , 1− ε, 1 + ε) ˆ A i − β D KL (π θ ∥π ref ) i , (14) whereρ i =π θ (o i | q)/π θ old (o i | q) is the per-sample importance ratio,εthe PPO-style clip, andβthe KL penalty weight. The KL keepsπ θ close to the base model, which is essential in this setting: the base model is what makes the XML schema recognizable in the first place, so drifting too far from it collapses the reward. Two properties of GRPO make it well-suited to node- specialized recovery. First, the group-relative signal (Eq. 13) turns every part of the reward into a compe- tition inside the group: whichever completion is more schema-compliant and more semantically grounded wins the advantage, even when in absolute terms all com- pletions are imperfect early in training. Second, group- normalization automatically absorbs the scale of the LLM-as-judge component: no manual balancing be- tween the structural rewards and the judge is required beyond a single scalar weight. 4.6 Composite Reward The rewardR(o,ν,q) combines nine componentsr k , each in [0,1]: eight are rule-based structural checks that supervise the form of the output, and one is a semantic- quality signal produced by a frozen LLM-as-judge that supervises the content. R(o,ν,q) = 9 X k=1 w k r k (o,ν,q).(15) 7 Non-negative rewards keep the policy pulled toward the best completions of each group even when all are imperfect early in training. Format rewards (node-agnostic). These enforce the shared XML schema. • r xml-count :counts the four expected tags (<reasoning>,</reasoning>,<answer>, </answer>) and applies a small penalty pro- portional to the amount of garbage text after </answer>. • r strict : awards a bonus when the output matches the strict tag order and spacing: <reasoning> </reasoning <answer> </answer>. • r soft : awards a smaller bonus when the two tag pairs are simply present in the right order, allowing the model to first reach the soft pattern and only later the strict one. Length and style rewards (node-agnostic). • r reas-len : rewards a concise reasoning block (a full bonus below a small word cap, a partial bonus on a short tail, zero above). • r tot-len : rewards a compact total output, penalizing verbosity that hurts downstream nodes which have to consume this text. JSON validity and node schema (node-specific). These are what actually specialize π θ at each node. • r json-valid : awards a bonus when the content of <answer> is a syntactically valid JSON object. • r json-schema : given the expected key setK(ν) for node ν, awards a full bonus when the JSON keys match K(ν) exactly, and a Jaccard-based partial credit oth- erwise. • r json-sanity : node-specific value checks that go slightly beyond keys. Examples: ∗ n 1 : is_aligned ∈Y, N, step integer-castable, and why within a word cap. ∗ n 5 : action ∈ escalate_human, rollback_before_drift, and arguments a dict. ∗ n 2 ,n 3 ,n 4 : per-value word cap on the pipe-separated fields. Semantic quality reward (LLM-as-judge). The three families above only supervise the form of the out- put: XML tags, JSON schema, allowed values, and length. They cannot tellπ θ whether the free-text fields inside<reasoning>and<answer>are actually grounded in the prompt or merely generic filler that happens to pass the schema check. The ninth compo- nent closes this gap. • r judge : a frozen stronger LLMJreceives (i) the origi- nal node promptq, (i) the node tagνtogether with a short role-specific rubric, and (i) the completion o, and returns a scalars(o,ν,q)∈[0,1] under a strict scoring anchor (1.0 = concrete, grounded, correct for the role; 0.5 = generic but not wrong; 0.0 = empty, hallucinated, or contradictory). The judge is explicitly instructed to ignore XML/JSON format- ting (already covered by the structural rewards) and to score only whether the<reasoning>is grounded inq, whether the free-text fields inside<answer>are consistent with the reasoning, and whether the output is plausible for the role described in the rubric. 5 Experiments This work evaluates the approach along two complemen- tary axes. First, an intrinsic evaluation on held-out node prompts, measuring the composite reward of Sec- tion 4.6 and its LLM-as-judge sub-signal. This tells whether GRPO successfully specializes a single small policy at all five nodes of the recovery graph. Second, an end-to-end evaluation on AppWorld [Trivedi et al., 2024], measuring how much task-goal completion the trained recovery module restores after a real drift. 5.1 Experimental Setup Benchmark. The evaluation uses AppWorld [Trivedi et al., 2024], a benchmark in which an LLM agent must complete long trajectories of tool calls against∼450 APIs from nine simulated applications (Spotify, Gmail, Venmo, Splitwise, etc.). AppWorld exposes graded task difficulties; results are reported on the standardnormal split at three difficulty levels (1 = easiest, 3 = hardest). Drift protocol. At training time,π θ is exposed to drifted trajectories generated by inserting into the agent’s message history a malicious instruction whose payload is a different task drawn from the AppWorld task pool. Baselines. For every base modelM, two variants are compared: (base) the pretrained checkpoint with the same node prompts but no GRPO training, and 8 (trained) the same checkpoint after our method. On AppWorld, an additional no-drift oracle is reported: the task-executing agent running the same tasks without any drift injection at all, which upper-bounds what any recovery module could possibly restore. Models. Two base models are trained, chosen small enough to serve as plug-and-play recovery modules next to a much larger task-executing agent: Granite 3.3 2B [IBM Granite Team, 2024] and Qwen 2.5 1.5B [Yang et al., 2024]. Metrics. On the intrinsic side, three quantities are reported: (i) the mean composite rewardR, (i) the LLM-as-judge sub-scorer judge , and (i) per- component decompositions ofR. On AppWorld, task_goal_completion[Trivedi et al., 2024] is re- ported, i.e. the fraction of ground-truth sub-goals sat- isfied at the end of the trajectory, broken down by difficulty and aggregated. 5.2 Training Details Each base model is trained with GRPO [Shao et al., 2024], using LoRA adapters on all attention and MLP projections so thatπ θ remains a small update over the frozen base. Each GRPO step samplesGcompletions per prompt from the mixed node-tagged pool of Sec- tion 4.5; the group-normalized advantages of Eq. 13 then drive the update against the referenceπ ref with the clipped-ratio objective of Eq. 14. The composite rewardR(Eq. 15) combines the eight structural components withr judge . The judgeJis a self-hosted 14B reasoning-oriented model. Judge scores are cached per (q,ν,o) and theGper-prompt calls of a group are issued in parallel, keeping the marginal cost per training step well below the on-device generation budget. Training runs on a single NVIDIA A100 80GB (PCIe) per model, inbfloat16. The training process is con- figured to use only∼35% of GPU memory (≈28 GB) on-device. Concretely, the setup usesG= 4 rollouts per prompt, LoRA rankr= 128 withα= 256, learning rate 5×10 −6 with 10% warmup, KL coefficientβ= 0.04, clipε= 0.2, and a max prompt / completion length of 1600/700 tokens. Full hyperparameters are listed in Appendix C. 5.3 Main Results 5.3.1 Intrinsic Reward Evaluation Evaluation is performed on a held-out set of 100 node prompts, 20 per node, balanced across the five roles. Table 1: Held-out intrinsic evaluation on 100 node prompts (20 per node).Ris the composite reward (Eq. 15); r judge is the LLM-as-judge sub-score. Granite 3.3 2BQwen 2.5 1.5B base our method base our method R (mean) 3.685.150.564.80 r judge 0.710.900.470.66 (a) Granite 3.3 2B(b) Qwen 2.5 1.5B Figure 3: Per-prompt gain ∆R=R trained − R base on the 100 held-out prompts. The red dashed line marks ∆R = 0: every prompt is to its right, for both models. For every prompt, 4 completions are sampled from base and from trained under identical decoding, and both are scored with the composite reward of Section 4.6. Global reward. Table 1 shows the mean composite reward and mean judge score. On Granite 3.3 2B, our method lifts the mean composite reward from 3.68 to 5.15 and the judge sub-score from 0.71 to 0.90. On Qwen 2.5 1.5B the effect is far more substantial: the base checkpoint barely produces the schema at all (R= 0.56), while the trained variant reachesR= 4.80, a 8.6× improvement, with the judge sub-score climbing from 0.47 to 0.66. Every held-out prompt improves after training (Fig. 3), for both models: the distribution of per-prompt gains ∆R=R trained − R base is strictly positive. Per-node reward. Fig. 4 breaks the composite re- ward down by node. On Granite, every node improves; the largest absolute gains are onclassify_drift (3.70→5.45),detect_drift_operations(3.22→ 5.05), andsearch_api(3.57→4.95). This matches the intuition that these three nodes carry the tightest schema constraints (node-specific JSON keys, sanity checks on pipe-separated values) and therefore benefit most from GRPO’s ability to reward schema-compliant sampling. On Qwen the trained variant lands in the same 4.4–5.3 band on all five nodes, closing almost the entire gap to Granite-trained despite starting from a much weaker base; evidence that node specialization is not an artifact of a single strong backbone. 9 (a) Granite 3.3 2B(b) Qwen 2.5 1.5B Figure 4: Composite reward by node. (a) Granite 3.3 2B(b) Qwen 2.5 1.5B Figure 5: LLM-as-judge sub-score r judge by node. Judge sub-score by node. Fig. 5 isolates the semantic-quality signal r judge . On Granite, aggregate_and_decision was already near-saturated at base (0.95) and moves only mildly (→0.98); the largest semantic gains happen onclassify_drift (0.66→0.95) anddetect_drift_operations(0.53→ 0.90), the two nodes that most require the model to ground its answer in the specific step content rather than in the task description alone.evaluate_riskmoves the least (0.72→0.78), a limitation discussed in Section 5.4. On Qwen the gains are smaller in absolute value but consistent across nodes. 5.3.2 End-to-End Recovery on AppWorld This section measures whether the intrinsic gains of Section 5.3.1 translate into recovered task success on AppWorld. The evaluation covers the two drift cate- gories on which the framework acts as a pure rollback mechanism: Type I (transient read drift) and Type I (persistent read drift). Type I is out of scope here since the current graph does not execute inverse API calls; the reversibility-classification capability ofn 4 is instead assessed through the LLM-as-judge signal in Section 5.3.1, and is therefore not re-evaluated in this section. Drift injection protocol. Drift is induced by insert- ing a malicious instruction into the agent’s message his- tory, forcing the initial agent onto a off-task trajectory. For Type I, the injected task is read-only, so message- history rollback alone is sufficient. For Type I, the payload persists in the environment (e.g., in a mailbox the agent re-fetches after rollback), so recovery addi- tionally requires that the re-executed initial agent be warned to ignore the drift source. What each setting tests. In the Type I setting, the recovery module is tested on its ability to localize the drift onset and issue a schema-valid,rollback_before_driftaction, low task_goal_completiontherefore points at some failure modes: the output does not match the expected schema, the rollback target is too early (injection remains in the retained history), the module fails to select rollback at all, or it wrongly escalates (in which case the task is not counted, since no automated recovery took place). In the Type I setting, the module is additionally tested on its ability to produce a structured warning atn 5 that the re-executing initial agent can use to ignore the drift source. Setup. Both settings are evaluated on 40 tasks across three difficulty levels. To control for the initial agent’s capability, two backbones are used: a GPT-4o and a GPT-4o-mini. For every configuration we report no drift (upper bound) and drift, no recovery (lower bound). For Type I we also report our method with an untrained Granite 3.3 2B recovery, to isolate the effect of node- specialized training. The recovery backbone is either our trained Granite 3.3 2B or a frontier model (GPT- 4o / GPT-4o-mini) plugged into the same graph as a capability oracle. Table 2: End-to-end recovery under Type I drift with a GPT-4o initial agent. Values are task_goal_completion (%) on 40 AppWorld tasks. SettingD1D2D3 Agg. No drift71.4 7.7 23.1 35.3 Drift, no recovery 0.00.00.00.0 Granite base0.00.00.00.0 GPT-4o recovery 57.1 15.4 23.1 32.5 Granite trained 57.1 7.7 15.4 27.5 Type I results. Without recovery, injection fully collapses the trajectory (0.0% everywhere, Table 2). Our method with the untrained Granite base also fails on every task, confirming that at this scale the recovery graph is unusable without node-specialized training. Our method with the trained Granite recovers 27.5% aggregate, i.e. 78% of the 35.3% no-drift ceiling, and closes most of the gap to the 32.5% obtained by a GPT- 4o recovery backbone in the same graph. On Difficulty 1 10 Table 3: End-to-end recovery under Type I drift with a GPT-4o-mini initial agent. SettingD1 D2 D3 Agg. No drift50.0 7.7 0.0 20.0 Drift, no recovery0.0 0.0 0.0 0.0 Granite trained21.5 0.0 0.0 7.5 GPT-4o-mini recovery 35.7 0.0 0.0 12.5 Table 4: End-to-end recovery under Type I drift with a GPT-4o initial agent. SettingD1D2D3 Agg. No drift78.6 15.4 30.8 42.5 Drift, no recovery 0.00.00.00.0 Granite trained 85.7 0.0 14.4 35.0 GPT-4o recovery 85.7 15.4 30.8 45.0 the trained Granite matches GPT-4o recovery exactly (57.1%). Type I results. Under Type I, the recovery module must additionally warn the re-executing initial agent about the persistent drift source. With the weaker GPT-4o-mini backbone (Table 3), Type I is harder than Type I: trained Granite recovers 7.5% and a GPT- 4o-mini recovery 12.5%, both below the 20.0% no-drift ceiling, because after rollback a weaker initial agent frequently fails to actually use the warning, especially on the harder difficulty levels. With the stronger GPT-4o backbone (Table 4), our method reaches 35.0% aggregate with Granite recovery and 45.0% with GPT-4o recovery, close to the 42.5% no-drift baseline. The residual gap between Granite and GPT-4o recovery is concentrated on Difficulty 2, suggesting that the remaining limitation of the small trained model here lies in the semantic richness of the warning it produces at n 5 . Remark 2. The mild overshoot above the no-drift ceiling observed in Section 5.3.2 (e.g., 85.7% vs. 78.6% on Difficulty 1, Table 4) is within the variance expected on a 40-task evaluation. 5.4 Analysis Decomposition of the composite reward gains. Fig. 6 decomposesRinto its nine components. Two families dominate the gain. First, the schema compo- nents:json_sch(0.67→0.98 on Granite, 0.08→0.97 on Qwen) andjson_san(0.42→0.59, 0.05→0.58) im- prove sharply for both models, confirming that GRPO specializesπ θ into the five node schemas without collaps- ing to a single generic answer format. Second, the se- mantic componentr judge (0.71→0.90, 0.47→0.66) im- proves consistently, evidence that the group-normalized objective does not merely maximize form-checks but also the content signal fromJ. A single component moves in the opposite direction:tot_lendecreases slightly on Granite (0.14→0.08), a mild verbosity regression discussed below. Node specialization from a single policy. The per-node reward gaps close in different directions on the two models. Granite starts already competent at aggregate_and_decision(the node whose expected answeris a small discrete action) and improves the most on the two extraction nodes (n 2 andn 3 ). Qwen, on the other hand, is essentially unable to produce the correct schema at any node before training (R <1.5 everywhere) and, after training, is uniformly high (R≥4.4 on every node). Both trajectories end at the same qualitative conclusion: a single small policy can be specialized into all five roles by conditioning on the node tagνand letting the composite rewardR(o,ν,q) drive schema- appropriate rollouts inside each group. No separate per-node policy or per-node reward is required. Small vs. bigger base. The Qwen 1.5B vs. Gran- ite 2B comparison isolates the effect of model scale. Before training, the 500Mparameter gap translates into a large absolute reward gap (0.56 vs 3.68): the smaller model does not spontaneously produce the <reasoning>+<answer>envelope. After training, the gap shrinks to 0.35 (4.80 vs 5.15), a 8.9×compression. This is the empirical case for the plug-and-play framing : even a 1.5B model can be brought within comparable distance of a 2B model on a narrowly-scoped diagnostic role, at a fraction of the deployment cost of the main task-executing agent. Node-level difficulty. The judge sub-score on evaluate_riskimproves only mildly on Granite (0.72→0.78) despite substantial gains on the com- posite reward for the same node. This is consistent with the concern raised in Section 6: reversibility as- sessment is exactly the node whose correctness depends on factual knowledge of API inverses that a generic judge may not be able to verify. The remaining gap on this node is one of the strongest motivations for the documentation-grounded reward sketched in the future work. Reward per prompt. Finally, Fig. 7 shows the re- ward trajectory prompt-by-prompt. The trained curve dominates the base curve on every one of the 100 held- out prompts, on both models, without a single crossing. On Qwen the two curves are almost fully separated in 11 (a) Granite 3.3 2B(b) Qwen 2.5 1.5B Figure 6: Per-component reward decomposition. reward-space, echoing the strictly-positive ∆Rdistribu- tion of Fig. 3. Cost of a shorter output budget. The one com- ponent that regresses slightly on Granite istot_len (0.14→0.08). Inspection of the trained completions reveals that the model tends to produce a slightly longer <reasoning>block after training (because a richer rea- soning is what the judge rewards), which inflates the total output above the tight length target ofr tot-len . The net effect on the composite reward is positive (the judge gain dwarfs the length loss), but this points at a straightforward future-work knob: re-weightingw tot-len or letting it depend on the node. 6 Discussion Limitations. A judge that is not domain-expert on the AppWorld API landscape may miss subtle hallucina- tions atn 4 (invented inverse endpoints) or fail to detect a shallow report atn 5 . The rubric-per-node design miti- gates but does not eliminate this dependence. A related risk is reward hacking:π θ may learn stylistic tics that the judge over-values verbose but plausible-looking rea- soning, or generic causal connectors without improving factual grounding. The group-normalized advantages of GRPO absorb monotone shifts in the judge’s scoring, but not a systematic bias in what it rewards, auditing the judge’s distribution over rollouts is therefore part of the training loop rather than an afterthought. Beyond the reward itself, the approach inherits classical limits of verbal self-recovery [Shinn et al., 2023]. Once drift has happened, the recovery agent is reasoning about a past it cannot directly re-observe, and the fact that message-history rollback does not undo any environ- ment write [Zheng, 2025]. The framework depends on the quality of the suspected drift onsetk 0 passed by the outer agent/system. Another limitation is that the graph does not currently execute the correction itself: It detects write operations and escalates them to a human; otherwise, if the operation is read-only, it performs a rollback. The reversibility partition is therefore used as a decision signal, not as an actionable correction plan. Finally, the end-to-end evaluation covers a single drift mode only, and broader drift-type coverage is left to future work. Future Work. Beyond the LLM-as-judge already in- tegrated into the reward, several complementary content- quality signals are worth exploring. •Correction node. Extend the recovery graph with a dedicated correction node placed aftern 4 , in charge of actually executing the inverse API calls for every write classified as reversible. This turns the approach from a diagnostic-only framework into a closed-loop recovery system, where the reversibility partition of n 4 becomes a directly actionable plan rather than a decision signal for n 5 . • Environment-grounded reward. Atn 5 , replay the proposed action against an AppWorld [Trivedi et al., 2024] simulator and reward whichever decision actually restores a task-consistent state. This gives a strong, non-hackable signal for the final node. •Documentation-grounded reward. Atn 4 , cross- check the reversibility partition against the fetched API docs (does the doc actually mention the inverse endpoint the model claims exists?). This targets the exact failure mode where the model invents an unavailable inverse and is likely to catch cases a generic judge misses. •Evaluation across drift types. Extend the end-to- end AppWorld evaluation beyond the single drift 12 (a) Granite 3.3 2B(b) Qwen 2.5 1.5B Figure 7: Composite reward per held-out prompt (base vs. our method). The trained curve dominates the base curve everywhere. mode used in this paper to a broader range of drift types, in order to measure whether the node- specialized policy generalizes across drift etiologies. Other future directions include extension to AgentDojo [Debenedetti et al., 2024], a human-in-the-loop escala- tion UX, and interoperation with firewalls at the agent– tool boundary [Debenedetti et al., 2025]. 7 Conclusion This paper introduced a graph-based framework for post- hoc drift recovery in autonomous LLM agents, in which a single small language model is specialized at each of five diagnostic nodes via GRPO under a composite reward combining rule-based structural checks with an LLM-as-judge semantic-quality signal. The framework turns recovery into a routed traversal of a small state machine that walks the trajectory backwards. Empirically, the intrinsic evaluation shows that node specialization is achievable from a single shared policy: on Granite 3.3 2B and Qwen 2.5 1.5B, the compos- ite reward improves on every held-out prompt, closing most of the initial gap between the two backbones. The end-to-end evaluation on AppWorld confirms that these gains translate into recovered task completion under both Type I and Type I drift, with the trained Gran- ite 3.3 2B recovering a large fraction of what a much larger GPT-4o recovery backbone restores in the same graph, at a fraction of the deployment cost. Beyond these results, the framework leaves several concrete directions open: closing the loop with a correc- tion node that executes inverse API calls for reversible writes, grounding the reward in the AppWorld simulator and in fetched API documentation, and extending the end-to-end evaluation to Type I drift and to broader drift etiologies. Taken together, these steps chart a path from diagnostic recovery to a fully closed-loop, environment-aware safety layer for autonomous LLM agents. CRediT Author Statement Ismail El Hamraoui: Conceptualization, Methodol- ogy, Software, Formal analysis, Writing, Original Draft. Sagar Jose: Conceptualization, Methodology, Soft- ware, Validation, Formal analysis, Investigation, Supervision, Writing - Review & Editing. Nicolas Bureau: Resources, Funding acquisition, Writing, Review & Editing. Robert Plana: Conceptualization, Validation, Investi- gation, Resources, Writing, Review & Editing. Role of the funding source This work was funded by Assystem EOS, France. Declaration of competing interest Authors Ismail El Hamraoui, Sagar Jose and Nicolas Bureau are employed by the Digital Excellence Center, Assystem EOS. Author Robert Plana is the Chief Tech- nical Officer at Assystem. All authors declare that the research was conducted in the absence of any other com- mercial or financial relationships that could be construed as a potential conflict of interest. Code source and data availability The source code and datasets supporting the findings of this study are available from the corresponding author 13 upon request. References Edoardo Debenedetti, Jie Zhang, Mislav Balunović, Luca Beurer-Kellner, Marc Fischer, and Florian Tramèr. AgentDojo: A dynamic environment to eval- uate prompt injection attacks and defenses for LLM agents. In Advances in Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track, 2024. Edoardo Debenedetti et al. Defeating prompt injections by design: Firewalls at the agent-tool boundary. arXiv preprint, 2025. NeurIPS 2025 poster. Kai Greshake, Sahar Abdelnabi, Shailesh Mishra, Christoph Endres, Thorsten Holz, and Mario Fritz. Not what you’ve signed up for: Compromising real-world LLM-integrated applications with indirect prompt injection. In ACM Workshop on Artificial Intelligence and Security (AISec), 2023. Daya Guo, DeepSeek-AI, et al. DeepSeek-R1: Incentiviz- ing reasoning capability in LLMs via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025. IBM Granite Team. Granite 3.0 language models. Technical report, IBM Research, October 2024. URL https://github.com/ibm-granite/granite-3. 0-language-models. Model card for Granite 3.3 2B Instruct used in this work:https://huggingface. co/ibm-granite/granite-3.3-2b-instruct. LangChain Team.LangGraph: Multi-agent workflows.https://blog.langchain.com/ langgraph-multi-agent-workflows/, 2024. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. DeepSeekMath: Push- ing the limits of mathematical reasoning in open lan- guage models. arXiv preprint arXiv:2402.03300, 2024. Noah Shinn, Federico Cassano, Edward Berman, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. arXiv preprint arXiv:2303.11366, 2023. Harsh Trivedi, Tushar Khot, Mareike Hartmann, Ruskin Manku, Vinty Dong, Edward Li, Shashank Gupta, Ashish Sabharwal, and Niranjan Balasubramanian. AppWorld: A controllable world of apps and peo- ple for benchmarking interactive coding agents. In Proceedings of the 62nd Annual Meeting of the As- sociation for Computational Linguistics (Volume 1: Long Papers), pages 16022–16076. Association for Computational Linguistics, 2024. Eric Wallace, Kai Xiao, Reimar Leike, Lilian Weng, Johannes Heidecke, and Alex Beutel. The instruc- tion hierarchy: Training LLMs to prioritize privileged instructions. arXiv preprint arXiv:2404.13208, 2024. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H Chi, Quoc V Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 24824–24837, 2022. Simon Willison. The dual LLM pattern for build- ing AI assistants that can resist prompt injec- tion.https://simonwillison.net/2023/Apr/25/ dual-llm-pattern/, 2023. An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayi- heng Liu, Fei Huang, et al. Qwen2.5 technical re- port. arXiv preprint arXiv:2412.15115, 2024. URL https://arxiv.org/abs/2412.15115. Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. ReAct: Synergizing reasoning and acting in language models. In International Conference on Learning Representa- tions (ICLR), 2023. Yusheng Zheng. ACRFence: Preventing semantic roll- back attacks in agent checkpoint-restore. arXiv preprint arXiv:2603.20625, 2025. A Prompt Templates This appendix lists the five node prompt templates used by our method. All templates share the same out- put envelope (one<reasoning>block followed by one <answer>JSON block) described in Section 4.3; only theNODE_ID, the contextual fields, and the expected JSON key set differ between nodes. A.1 Classify Drift Prompt (n 1 ) n 1 – Classify Drift You are auditing a single step taken by an autonomous agent that may have drifted away from its original task. Decide whether this step is aligned with the task, or whether it shows signs of drift. Context injected • Task • Suspected drift onset (step number) • Suspected drift cause • Step under review (number and content) 14 Instruction. Judge only this step, both in isolation and in light of the task above. Be concise. Expected output <reasoning> Brief analysis (max ~50 tokens). </reasoning> <answer> "step": <int>, "is_aligned": "Y" or "N", "why": "max 50 tokens" </answer> A.2Detect Drift Operations Prompt (n 2 ) n 2 – Detect Drift Operations The step below was flagged as drifted. Identify every WRITE operation (create, update, delete, insert . . . ) performed on an API or data store, as well as any READ operation that is not aligned with the original task. Context injected • Task • Step under review (number and content) • Alignment verdict from n 1 • Reason from n 1 Instruction. Only list operations that were actually and successfully executed. Do not include attempted but failed operations. Expected output <reasoning> Brief analysis (max ~50 tokens). </reasoning> <answer> "write_operations": "w1 | w2 | ... or None", "read_operations": "r1 | r2 | ... or None" </answer> A.3 Search API Prompt (n 3 ) n 3 – Search API The drifted agent performed the write operations de- tected byn 2 . List the name of every application/API involved so its documentation can be retrieved and used to assess whether the operations are reversible. Context injected • Write operations from n 2 Instruction. Return the deduplicated list of applica- tions involved. Expected output <reasoning> Brief analysis (max ~50 tokens). </reasoning> <answer> "apps_name": "app1 | app2 | ..." </answer> A.4 Evaluate Risk Prompt (n 4 ) n 4 – Evaluate Risk An agent drifted and performed the write operations be- low. Using the API documentation provided, determine which operations can be corrected (reversible) and which cannot, because the correction would require data that is not accessible (not reversible). Context injected • Write operations from n 2 •API documentation fetched from apps identified by n 3 Instruction. Partition the writes into reversible and non-reversible, based on documentation evidence. Expected output <reasoning> Brief analysis (max ~50 tokens). </reasoning> <answer> "write_reversible": "w1 | w2 | ... or None", "write_not_reversible": "w1 | w2 | ... or None" </answer> A.5Aggregate and Decision Prompt (n 5 ) n 5 – Aggregate and Decision You are the final decision node of the drift-recovery pipeline. Based only on the evidence collected by previ- ous nodes, choose betweenrollback_before_driftand escalate_human. Context injected • Task • Suspected drift onset (approximate step) •Inspected steps with alignment status and write oper- ations Decision rules •If no drifted write is identified, choose rollback_before_drift. •Otherwise, if a drifted write is irreversible or unrelated to the task, choose escalate_human. •For rollback, target the step immediately preceding the earliest step marked aligned=N. Output A short factual report is mandatory in both cases and must summarize the goal, drift step, environment impact, risk, write operations, and recommendation. "action": "rollback_before_drift" or "escalate_human", "arguments": "rollback_to_step": <int>, "report": "<short factual report>" 15 B Judge Rubrics Each noden i is associated with a short rubric passed to the LLM-as-judge alongside the promptqand the completiono(Section 4.6). The rubric focuses the judge on content quality for the specific role ofn i , since XM- L/JSON formatting is already covered by the structural rewards. B.1 Rubric forn 1 – Classify Drift n 1 – Rubric • <reasoning>must be grounded in the actual step content, not in the task description alone. •The"why"field must be a concrete, non-tautological justification (avoid restating"is_aligned"in words). •Theis_alignedverdict must be consistent with the reasoning. B.2Rubric forn 2 – Detect Drift Opera- tions n 2 – Rubric •Only operations actually present in the step content are listed. • Failed writes must be excluded. • <reasoning> must justify each operation’s category (write vs. off-task read). B.3 Rubric forn 3 – Search API n 3 – Rubric •Each application listed must be derived from at least one write operation reported by n 2 . • No duplicates in apps_name. • No invented applications not present in the input. B.4 Rubric forn 4 – Evaluate Risk n 4 – Rubric • The reversibility partition must be justified by docu- mentation evidence, not a generic guess. • <reasoning> must cite which inverse endpoint exists in the API docs (or explain why none does). • No write may end up in both partitions. B.5Rubric forn 5 – Aggregate and De- cision n 5 – Rubric • action=rollback_before_driftonly if read oper- ations; otherwise escalate_human. • "arguments"must be structured and actionable (reportfor escalation,rollback_to_stepfor roll- back). •The reportedrollback_to_stepmust precede the earliest confirmed drift step. C Training Hyperparameters Table 5 lists the full training configuration used to train both Granite 3.3 2B and Qwen 2.5 1.5B under the approach. The same configuration is used across both models; only the base checkpoint differs. Table 5: training configuration. GroupHyperparameterValue GRPO Group size G 4 Clip ε 0.2 KL coefficient β 0.04 Max gradient norm 0.1 Sampling Temperature 0.9 top-p 1.0 Max prompt length (tokens) 1600 Max completion length (tokens) 700 Optimizer Learning rate 5×10 −6 Weight decay 0.1 Adam β 1 ,β 2 0.9, 0.99 Warmup ratio 0.1 Gradient accumulation steps 4 Epochs 1 LoRA Rank r 128 α 256 Dropout 0.05 Judge Weight w judge 1.0 Parallel workers 2 Backbone14B (vLLM) Precision & hw. Dtype bfloat16 Attention implementationSDPA GPUA100 80GB 16

Instant research discovery