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

1 Introduction

• Agents increasingly call APIs, send emails, transfer money, and modify databases, making mistakes more costly as they act on real systems.The introduction names ReAct [Yao et al., 2023] as popularizing this pattern and AppWorld [Trivedi et al., 2024] as a benchmark of long tool-call trajectories.Source: §1 · p. 1
• Behavioral drift begins when an agent loses sight of its initial goal and can leave persistent world changes that memory reset cannot undo.The passage identifies malicious instructions in webpages or email and an agent's over-extension of a vague sub-goal as possible causes.Source: §1 · p. 1
• Its recovery procedure locates drift onset, inventories committed write operations, assesses inverse-API undoability, then chooses message-history rollback or human escalation.In some cases, the original task cannot resume until a correction action is taken or a human repairs the environment.Source: §1 · pp. 1–2
• The paper frames drift recovery as a five-node diagnostic state machine that routes backward through the trajectory, using one small LLM specialized across roles with role-conditioned prompts and shared GRPO.The model emits a <reasoning> block followed by an <answer> JSON block, and the traversal seeks an aligned streak of steps.Source: §1 · p. 2

2.1 Attacks and Defenses on LLM Agents

• Prompt injection can hijack tool-using agents [Greshake et al., 2023], while common defenses target prevention; this work instead assumes drift may occur and assesses actions already taken [Zheng, 2025].AgentDojo [Debenedetti et al., 2024] benchmarks injection success and effects on the original task; proposed defenses isolate untrusted content, teach instruction hierarchies [Wallace et al., 2024], or firewall tools.Source: §2.1 · p. 2

2.2 Graph-based Agents

• Unlike ReAct-style loops, the five-node graph assigns distinct diagnostic roles and reconstructs past actions in a drifted trajectory instead of advancing the task.The nodes classify, extract, look up, evaluate, and decide.Source: §2.2 · p. 2

2.3 Reinforcement learning for LLMs

• GRPO [Shao et al., 2024] suits training because parseable outputs provide verifiable rewards, group-relative feedback remains useful early, and omitting a value network enables single-GPU training.GRPO drops the value network and computes advantages relative to same-prompt rollouts; rule-based verifiable rewards have previously elicited long reasoning without human preference data [Guo et al., 2025].Source: §2.3 · p. 2

2.4 Structured Reasoning

• A fixed <reasoning>+<answer> schema makes each node’s output a typed message for the next, avoiding the reparsing required by free-form reasoning [Wei et al., 2022; Shinn et al., 2023].Chain-of-Thought uses intermediate steps, and Reflexion adds a natural-language critique of a past attempt; this pipeline instead fixes the output structure.Source: §2.4 · p. 2

3.1 Notation and Setup

• Read operations observe external entities without changing them, while write operations mutate them; entities include APIs, databases, file systems, applications, and other tool-exposed services.Each step’s interaction is decomposed into read and write operation sets.Source: §3.1 · p. 3

3.2 Ideal Executions and Task Success

• Because tasks can have multiple valid trajectories, success allows semantic equivalence under independent-read reordering or API-path substitutions that preserve the final environment state.AppWorld’s task_goal_completion measures this form of equivalence [Trivedi et al., 2024].Source: §3.2 · p. 3

3.3 Behavioral Drift

• Drift begins at step d when execution matches a task-faithful reference before d but cannot maintain one from d onward; drifted and task-consistent steps may then interleave.The work later assumes monotone drift, ruling out this interleaved pattern.Source: §3.3 · p. 3, §3.5 · p. 4

3.4 Drift Taxonomy

• Type I transient read drift can be recovered by message-history rollback, whereas Type II persistent read drift requires recovery memory to avoid re-exposure to the same trigger [Shinn et al., 2023].Type I’s source is transient; Type II’s persists in the environment or task inputs, so naive rollback is insufficient.Source: §3.4 · p. 4
• Type III drift includes a write, so recovery requires repairing or compensating documented reversible modifications and rolling back and re-executing; irreversible writes require human escalation [Zheng, 2025].Rollback of agent state alone does not reset the environment.Source: §3.4 · p. 4

3.5 Scope of This Work

• The work focuses on Type I, uses drift-aware memory for Type II, and escalates Type III once any write is committed rather than automatically compensating it; inverse-API correction is future work.This scope assumes drift is monotone after onset.Source: §3.5 · p. 4

3.6 Recovery Acceptance Criteria

• Recovery criteria require drift-robust, autonomous, state-aware diagnosis and reversibility-sensitive correction or human escalation, with reversible actions distinguished from irreversible ones.Diagnosis must identify root cause and onset using information available at recovery time, without an oracle or ground-truth trace.Source: §3.6 · p. 4
• The framework should be plug-and-play without retraining the main agent, resource-efficient, and preventive by surfacing incident structure to help guard against recurring drift.The main task-executing model is typically large, expensive, and shared across deployments.Source: §3.6 · p. 5

4.1 Overview

• The recovery graph has five nodes, all served by the same policy through role-specialized prompts.Source: §4.1 · p. 5, §4.2 · p. 5

4.2 Recovery Graph Architecture

• Classify Drift (n1) uses an external detector’s suspected onset to judge whether a single step is task-aligned or drifted, rather than detecting drift autonomously.Source: §4.2 · p. 5
• When n1 flags drift, Detect Operations (n2) lists successfully executed writes and off-task reads, excluding failed writes and task-aligned reads.Source: §4.2 · p. 5
• Search Documentation (n3) deduplicates the applications involved in the writes, and Evaluate Risk (n4) partitions those writes by reversibility using fetched API documentation.Source: §4.2 · p. 5

5. Aggregate

• The final node chooses rollback_before_drift with a target step or escalate_human with a structured incident report.It makes this decision from the full per-step diagnosis.Source: §5 · p. 5

4.3 Structured XML Reasoning

• Each node requires exactly one <reasoning> block followed by one <answer> block containing valid JSON with node-specific keys.The specified key sets include alignment diagnosis at n1, operation detection at n2, reversibility assessment at n4, and action plus arguments at n5.Source: §4.3 · p. 6

4.4 Backward Walk and Aligned-Streak Termination

• The graph walks backward after misalignment, processes detected writes, and terminates only after K consecutive aligned verdicts.A single aligned step is insufficient evidence that drift has ended; the short streak is used as a robust, inexpensive proxy.Source: §4.4 · p. 6

4.5 Node-Specialized RL Training with GRPO

• One shared policy is specialized across all five graph nodes with GRPO, which compares groups of structured completions without a value network.Group-normalized advantages let better schema-compliant and semantically grounded completions win within each group, while normalization absorbs the judge signal’s scale without manual balancing.Source: §4.5 · p. 7, Composite reward · p. 7

Composite reward.

• Group-relative advantages avoid a value function, and a KL penalty keeps the policy close to its pretrained base model.The method also uses group normalization to absorb the judge component’s scale without manual balancing against structural rewards.Source: Composite reward · p. 7

4.6 Composite Reward

• The GRPO reward combines eight rule-based structural checks with one frozen LLM judge that evaluates semantic quality.The judge scores grounding and consistency in the reasoning and free-text answer fields, while ignoring formatting already evaluated by structural rewards.Source: §4.6 · pp. 7–8

5.1 Experimental Setup; 5.2 Training Details

• The study evaluates AppWorld recovery on the standard normal split at three difficulty levels, covering Type I transient read drift and Type II persistent read drift; Type III is excluded because the graph does not execute inverse API calls.Source: §5.1 · p. 8, §5.3.2 · p. 10

5.3.1 Intrinsic Reward Evaluation

• Granite’s mean composite reward rose from 3.68 to 5.15, while Qwen’s rose from 0.56 to 4.80; every held-out prompt improved for both models.Granite’s judge sub-score increased from 0.71 to 0.90, and Qwen’s from 0.47 to 0.66. Evaluation used 100 held-out node prompts, with 20 prompts per node.Source: Table 1 · p. 9, Figure 3 · p. 9, §5.3.1 · p. 9
• Composite reward improved at every Granite node, with the largest gains on classify_drift, detect_drift_operations, and search_api; trained Qwen reached 4.4–5.3 across all five nodes.Granite’s largest gains included classify_drift (3.70 → 5.45), detect_drift_operations (3.22 → 5.05), and search_api (3.57 → 4.95); these nodes have tighter schema constraints.Source: §5.3.1 · p. 9, Figure 4 · p. 10
• Granite’s judge sub-score rose most on classify_drift (0.66 → 0.95) and detect_drift_operations (0.53 → 0.90), but evaluate_risk improved only from 0.72 to 0.78.The paper attributes the evaluate_risk limitation to the factual API-inverse knowledge needed for reversibility assessment, which a generic judge may not verify; Qwen’s judge gains were smaller but consistent across nodes.Source: Figure 5 · p. 10, §5.3.1 · p. 10, §5.4 · p. 11

5.3.2 End-to-End Recovery on AppWorld

• On Type I drift, trained Granite recovered 27.5% of tasks, while Type II recovery reached 35.0% with a GPT-4o initial agent.For Type I, 27.5% is 78% of the 35.3% no-drift ceiling and approaches the 32.5% GPT-4o recovery result; untrained Granite and no recovery both scored 0.0%. For Type II with GPT-4o-mini, trained Granite recovered 7.5% against a 20.0% no-drift ceiling, while GPT-4o recovery reached 12.5%; with GPT-4o, the corresponding results were 35.0% for Granite and 45.0% for GPT-4o against a 42.5% no-drift baseline. The weaker initial agent often failed to use the warning needed to handle persistent drift.Source: Table 2 · p. 10, §5.3.2 · pp. 10–11, Table 3 · p. 11

5.4 Analysis

• A single node-conditioned policy reached high rewards across all five roles, and training narrowed the Qwen–Granite reward gap from 3.12 to 0.35.Qwen began below 1.5 reward at every node and reached at least 4.4 after training; the analysis concludes that separate per-node policies or rewards are unnecessary. The reported gap compression is 8.9×.Source: §5.4 · p. 11
• Schema rewards improved sharply, including json_sch from 0.67 to 0.98 on Granite and from 0.08 to 0.97 on Qwen, alongside gains in semantic judge scores.The analysis reports json_san increasing from 0.42 to 0.59 on Granite and from 0.05 to 0.58 on Qwen; Granite’s tot_len component decreased from 0.14 to 0.08 as richer reasoning lengthened outputs beyond the target.Source: §5.4 · pp. 11–12, Figure 6 · p. 12

Limitations

• A non-expert judge may miss subtle API hallucinations or shallow reports, while reward hacking can favor plausible style over factual grounding.GRPO can absorb monotone shifts in judge scores but not systematic reward bias, so auditing the judge’s rollout-score distribution remains part of training.Source: §6 · p. 12
• Recovery depends on the outer system’s suspected onset and cannot re-observe the past; rolling back message history does not undo environment writes.The framework also inherits limits of verbal self-recovery [Shinn et al., 2023].Source: §6 · p. 12
• The graph does not execute corrections: it escalates detected writes to a human, rolls back read-only operations, and evaluation covers only one drift mode.Reversibility currently guides the decision rather than providing an actionable correction plan.Source: §6 · p. 12

Future Work

• A proposed correction node after n4 would execute inverse API calls for writes classified as reversible, making the framework closed-loop.This would turn n4’s reversibility partition from a decision signal into an actionable recovery plan.Source: §6 · p. 12
• Proposed grounded rewards would test n5 decisions in an AppWorld simulator [Trivedi et al., 2024] and cross-check n4 reversibility claims against API documentation.The simulator would reward restoration of a task-consistent state; documentation checks target invented inverse endpoints that a generic judge may miss.Source: §6 · p. 12
• Future evaluation would test whether the node-specialized policy generalizes across drift etiologies by extending AppWorld beyond the paper’s single drift mode.The discussion also proposes extension to AgentDojo [Debenedetti et al., 2024], a human-in-the-loop escalation interface, and interoperability with agent-boundary firewalls.Source: §6 · pp. 12–13

7 Conclusion

• The framework routes recovery through a small state machine traversing the trajectory backwards, specializing a single small language model at each of five diagnostic nodes via GRPO.Its composite reward combines rule-based structural checks with an LLM-as-judge semantic-quality signal.Source: §7 · p. 13
• AppWorld evaluation shows trained Granite 3.3 2B recovers a large fraction of the task completion restored by GPT-4o under both Type I and Type II drift, at lower deployment cost.Intrinsic evaluation also finds that the composite reward improves performance on every held-out prompt for Granite 3.3 2B and Qwen 2.5 1.5B, closing most of their initial gap.Source: §7 · p. 13
• Future directions include adding a correction node for inverse API calls, grounding rewards in the simulator and fetched API documentation, and evaluating Type III drift and broader etiologies.Together, these extensions point toward a fully closed-loop, environment-aware safety layer for autonomous LLM agents.Source: §7 · p. 13

• El Hamraoui, Jose, and Bureau are employed by Assystem EOS’s Digital Excellence Center, Plana is Assystem’s Chief Technical Officer, and all authors declare no other potential conflicting commercial or financial relationships.Source: Declaration of competing interest · p. 13

A Prompt Templates

• All five templates use the same output envelope, with node-specific IDs, context fields, and JSON keys.The envelope is one <reasoning> block followed by one <answer> JSON block.Source: A Prompt Templates · p. 14

A.1 Classify Drift Prompt (n1)

• Node n1 judges whether a step is aligned with the task, using the task, suspected onset and cause, and step under review.It assesses only that step, in isolation and in light of the task, and returns an alignment verdict with a brief rationale.Source: §A.1 · pp. 14–15

A.2 Detect Drift Operations Prompt (n2)

• Node n2 identifies executed, successful writes and task-misaligned reads from a flagged step, using n1’s verdict and rationale.It excludes attempted but failed operations and returns separate write- and read-operation lists.Source: §A.2 · p. 15

A.3 Search API Prompt (n3)

• Node n3 deduplicates the applications or APIs involved in n2’s write operations so their documentation can be retrieved.Its output is a list of application names.Source: §A.3 · p. 15

A.4 Evaluate Risk Prompt (n4)

• Node n4 uses n2’s writes and n3’s API documentation to partition operations into reversible and non-reversible categories.The distinction depends on whether correction is possible with the available data, as supported by documentation evidence.Source: §A.4 · p. 15

A.5 Aggregate and Decision Prompt (n5)

• Node n5 chooses rollback or human escalation from prior evidence: no drifted write prompts rollback; an irreversible or unrelated drifted write prompts escalation.Rollback targets the step immediately before the earliest step marked aligned=N; either action requires a short factual report covering the goal, drift step, impact, risk, writes, and recommendation.Source: §A.5 · p. 15

B Judge Rubrics

• The judge evaluates role-specific content quality, with structural rewards separately handling XML/JSON formatting.Each node’s rubric is passed to the LLM-as-judge alongside the prompt and completion.Source: B Judge Rubrics · p. 16

B.1 Rubric for n1 – Classify Drift

• For drift classification, reasoning must use actual step content, the “why” must be concrete, and the alignment verdict must match the reasoning.Source: §B.1 · p. 16

B.2 Rubric for n2 – Detect Drift Operations

• For operation detection, list only operations in the step, exclude failed writes, and justify each operation’s write or off-task-read category.Source: §B.2 · p. 16

B.3 Rubric for n3 – Search API

• For API search, derive each application from a reported write operation, avoid duplicate names, and do not invent applications absent from the input.Source: §B.3 · p. 16

B.4 Rubric for n4 – Evaluate Risk

• For risk evaluation, justify reversibility using documentation, identify the inverse endpoint or its absence, and place each write in only one partition.Source: §B.4 · p. 16

B.5 Rubric for n5 – Aggregate and Decision

• For the final decision, choose rollback only for read operations and otherwise escalate; provide actionable arguments and set rollback before the earliest confirmed drift step.Source: §B.5 · p. 16

• Table 5 presents the full training configuration used to train Granite 3.3 2B and Qwen 2.5 1.5B.Source: C Training Hyperparameters · p. 16, Table 5 · p. 16
• Both models use the same training configuration; only their base checkpoints differ.Source: C Training Hyperparameters · p. 16