PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.14109v1cs.AIcs.LGcs.MA

TL;DR

Runtime drift in autonomous LLM agents lacks structured, step-level recovery, so this paper trains one small language model across a graph of specialized diagnostic roles. On AppWorld, the approach generally issues correct recovery decisions under Type I and Type II drift while reliably following the required schema and producing semantically appropriate fields.

  • Problem

    Existing drift defenses focus on inputs and lack a structured procedure for locating drift onset, identifying committed writes, assessing reversibility, and choosing recovery actions.

  • Method

    A graph-based recovery state machine specializes one shared small language model across diagnostic nodes using reinforcement learning and structured outputs with composite structural and semantic rewards.

  • Results

    On AppWorld, the trained Granite 3.3 2B recovers a large fraction of task completion restored by GPT-4o under both Type I and Type II drift, at lower deployment cost.

  • Takeaways & Limitations

    A single shared small policy can support structured, multi-role drift recovery while translating node-level training gains into recovered task completion.

  • Takeaways & Limitations

    Automated inverse API correction is deferred, and the evaluation uses a single drift mode rather than testing generalization across broader drift types.

Abstract

from arXiv · show

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.

1 Introduction

As LLM agents take consequential actions through external systems, behavioral drift can cause costly deviations from users’ goals. This paper introduces a graph-based recovery framework that routes a single small LLM through specialized diagnostic roles.

  • Motivation: LLM agents increasingly call APIs, send emails, transfer money, and modify databases, making mistakes more costly as actions affect real systems.AppWorld evaluates such agents through long tool-call trajectories across hundreds of APIs.
  • Problem: Behavioral drift begins when an agent loses sight of its initial goal, potentially after malicious content or self-directed overextension of a vague sub-goal.Drift differs from a crash or merely hallucinated answer, though hallucinations can trigger actions leading to drift.
  • Problem: Existing drift defenses focus on input prevention, leaving a post-drift procedure to locate onset, identify committed writes, assess reversibility, and choose recovery escalation.The procedure must walk the trajectory backwards and decide between message-history rollback and human escalation.
  • Contribution: A graph-based framework turns drift recovery into a routed diagnostic state machine with five nodes, each specializing a single small LLM for a recovery role.The framework addresses the brittleness of asking one LLM to perform all subtasks in one prompt and the schema limitations of free-form reasoning.
  • Contribution: One shared small LLM is specialized across multiple roles through role-conditioned prompts and a shared GRPO objective, while graph traversal searches backward for an aligned streak of steps.The method avoids training a separate policy per node.

2 Related Work

Prior work primarily prevents prompt-injection-driven drift or uses general agent loops, whereas this work diagnoses already-drifted trajectories with role-specific graph nodes. It also replaces free-form reasoning with typed node outputs and uses GRPO for verifiable, resource-efficient training.

  • Prompt-injection defenses: Prompt-injection defenses mainly prevent attacker-controlled content from steering tool-using agents by isolating untrusted text, imposing instruction hierarchies, or filtering tool access.Benchmarks such as AgentDojo measure attack success and task disruption.
  • Post-drift recovery: This work instead assumes drift has occurred and reconstructs the affected trajectory region, including environment writes, before deciding how to respond.Resetting message history cannot undo writes already made to the environment.
  • Graph-based diagnosis: Unlike ReAct-style loops and multi-node state machines that advance task execution, the proposed five-node graph assigns distinct diagnostic roles and reconstructs past behavior.Its roles are classify, extract, look up, evaluate, and decide.
  • Reinforcement learning: GRPO trains the recovery model using group-relative advantages without a value network, matching rewards that verify parsing and JSON-key correctness.The absence of a value network enables training on a single GPU.
  • Structured reasoning: The pipeline uses a fixed <reasoning>+<answer> schema so each node emits a typed message that the next node can read without reparsing free-form text.This is a structured alternative to free-form Chain-of-Thought and natural-language self-critique.

3 Problem Formulation

The problem formulation models autonomous-agent execution as environment-facing traces, defines drift relative to task-faithful trajectory sets, and distinguishes recovery requirements by drift type. It also specifies scope assumptions and acceptance criteria for a plug-and-play recovery module.

  • 3.1 Notation and Setup: Each execution step comprises read operations that leave external entities unchanged and write operations that mutate them, forming the agent’s environment-facing trace.External entities include APIs, databases, file systems, applications, and other tool-exposed services.
  • 3.2 Ideal Executions and Task Success: Task-faithful behavior is represented by a set of valid trajectories because executions may differ in independent-read order, tool-call decomposition, or API path.Task success requires equivalence to at least one reference trajectory under the chosen semantic equivalence relation.
  • 3.3 Behavioral Drift: Drift begins at onset d when execution matches a task-faithful reference before d but cannot maintain any such reference from d onward.Drift is defined relative to a task-faithful trajectory rather than a single canonical trace.
  • 3.3 Behavioral Drift: Post-onset behavior need not be uniformly off-task: drifted and task-consistent steps can be interleaved after the initial deviation.Recovery feasibility and cost depend on both suspected onset and intervening actions, especially irreversible writes.
  • 3.4 Drift Taxonomy: The taxonomy distinguishes transient read drift, persistent read drift, and write drift according to environment mutation and causal persistence.Transient read drift permits message-history rollback, whereas persistent read drift additionally requires retaining drift knowledge in recovery context.
  • 3.4 Drift Taxonomy: Write drift requires environment repair or compensation plus agent-state rollback and re-execution; any irreversible write makes purely automated recovery unsafe and requires human escalation.Resetting agent state alone cannot restore the prior environment state.
  • 3.5 Scope of This Work: The work primarily targets Type I drift, applies drift-aware memory for Type II, and authorizes rollback only when the drifted region contains reads exclusively.For Type III drift, the framework detects reversibility but does not automatically compensate environment modifications; future work proposes inverse API calls for reversible writes.
  • 3.6 Recovery Acceptance Criteria: The recovery framework must resist drift, diagnose its cause and exact onset without an oracle, reason about environment state, and integrate without retraining the main model.The formulation additionally assumes monotone drift after onset, excluding interleaved post-onset behavior.

4 Methodology

The methodology models drift recovery as a routed traversal of a five-node state machine, using one shared policy with role-specialized prompts at each node. The graph progressively classifies drift, identifies relevant operations, locates documentation, and evaluates write reversibility.

  • Recovery Graph: Recovery uses a five-node graph in which the same policy receives role-specialized prompts at each traversal stop.The approach treats recovery as a routed traversal of a small state machine.
  • Graph Nodes: Classify Drift decides whether each externally flagged step is task-aligned or drifted and provides a short reason.It refines an external detector’s suspected onset rather than detecting drift autonomously.
  • Graph Nodes: Detect Operations activates only for drifted steps and enumerates successful writes plus off-task reads, excluding failed writes and aligned reads.This node operates on the content of the flagged step.
  • Graph Nodes: Search Documentation deduplicates the applications involved in the detected writes so their API documentation can be fetched.The node outputs application identifiers such as spotify | gmail.
  • Graph Nodes: Evaluate Risk partitions the detected writes into write_reversible and write_not_reversible using the fetched API documentation.Its inputs are the writes and the corresponding documentation.

5. Aggregate

The recovery graph combines backward diagnosis, streak-based termination, and a final rollback-or-escalation decision. A shared model is specialized across node roles with GRPO using structural and semantic rewards for concise, schema-valid, grounded outputs.

  • Recovery graph: n5 converts the per-step diagnosis into either rollback_before_drift with a target step or escalate_human with a structured incident report.
  • Recovery graph: The controller walks backward after drift, routes through operation and risk nodes when writes occur, and terminates only after K consecutive aligned verdicts.A single aligned step immediately before a drifted region is treated as insufficient evidence that drift has ended.
  • Training: The same model πθ serves all five roles, while node-specific prompts and answer schemas specialize behavior through a shared GRPO rollout procedure.Each completion contains a reasoning block followed by node-specific answer JSON.
  • Training: The reward combines eight structural checks with one frozen LLM-as-judge signal, supervising both output form and semantic grounding.Structural checks cover XML format, length, JSON validity, node schema, and sanity constraints; the judge scores role-specific correctness and grounding.

5 Experiments

The experiments evaluate both intrinsic node specialization and end-to-end drift recovery on AppWorld. Training substantially improves structured rewards and recovery outcomes, while revealing weaker performance on risk assessment and Type II recovery.

  • Evaluation design: The evaluation combines held-out node-prompt rewards with end-to-end AppWorld task-goal completion across three difficulty levels.AppWorld contains long tool-call trajectories spanning approximately 450 APIs across nine simulated applications.
  • Intrinsic results: 8.6×: Qwen 2.5 1.5B’s mean composite reward rises from 0.56 to 4.80, while Granite 3.3 2B improves from 3.68 to 5.15.Judge sub-scores increase from 0.47 to 0.66 for Qwen and from 0.71 to 0.90 for Granite.
  • Limitations: Risk assessment remains the main semantic weakness, with Granite’s evaluate_risk judge sub-score improving only from 0.72 to 0.78.The limitation is attributed to factual API-inverse knowledge that a generic judge may not verify.
  • End-to-end recovery: 27.5%: trained Granite recovers aggregate Type I task-goal completion, reaching 78% of the 35.3% no-drift ceiling and matching GPT-4o recovery at 57.1% on Difficulty 1.Without recovery and with untrained Granite, completion is 0.0% everywhere.
  • End-to-end recovery: 7.5%: trained Granite recovers Type II completion with GPT-4o-mini as the initial agent, below the 20.0% no-drift ceiling because the weaker agent often fails to use the warning.The GPT-4o-mini recovery variant reaches 12.5% under the same setting.
  • Reward decomposition: Schema and semantic components drive the reward gains, including json_sch increases from 0.08 to 0.97 on Qwen and from 0.67 to 0.98 on Granite.The judge sub-score also improves consistently, while Granite’s tot_len decreases from 0.14 to 0.08.
  • Node specialization: After training, Qwen reaches R ≥4.4 on every node, whereas its untrained checkpoint has R <1.5 everywhere; Granite improves most on classify_drift and detect_drift_operations.The trained reward curve dominates the base curve on all 100 held-out prompts for both models.

6 Discussion

The framework remains limited by judge quality, dependence on the suspected drift onset, lack of executable correction, and evaluation on a single drift mode. Future work targets closed-loop correction, stronger environment- and documentation-grounded rewards, broader drift evaluation, and expanded integrations.

  • Limitations: Judge limitations include missed subtle hallucinations or shallow reports, reward hacking, and dependence on the suspected drift onset k0 supplied by the outer system.The rubric-per-node design mitigates but does not eliminate these risks.
  • Limitations: The graph does not execute corrections: it escalates write operations to a human and rolls back read-only operations, using reversibility only as a decision signal.The end-to-end evaluation also covers only one drift mode.
  • Future Work: A dedicated correction node after n4 could execute inverse API calls for reversible writes, turning the diagnostic framework into closed-loop recovery.This would make n4’s reversibility partition an actionable plan rather than only an input to n5.
  • Future Work: Environment-grounded reward would replay n5’s proposed action in an AppWorld simulator and reward decisions that restore a task-consistent state.Documentation-grounded reward at n4 would verify claimed inverse endpoints against fetched API documentation and target invented unavailable inverses.
  • Future Work: Future evaluation should cover broader drift types to test whether the node-specialized policy generalizes across drift etiologies.Other directions include AgentDojo, human-in-the-loop escalation UX, and interoperation with agent-boundary firewalls.

7 Conclusion

The paper presents a graph-based posthoc drift-recovery framework using one small language model specialized across diagnostic nodes with composite structural and semantic rewards. Evaluations support node specialization and recovered task completion, while identifying extensions toward closed-loop, environment-aware safety.

  • Framework: The framework routes autonomous-agent recovery through a backward-traversed state machine with one small model specialized at five diagnostic nodes using GRPO.Training combines rule-based structural checks with an LLM-as-judge semantic-quality signal.
  • Evaluation: On Granite 3.3 2B and Qwen 2.5 1.5B, composite reward improves on every held-out prompt, closing most of the initial backbone gap.This intrinsic evaluation indicates that node specialization is achievable from a single shared policy.
  • Evaluation: AppWorld end-to-end evaluation confirms that these gains translate into recovered task completion under both Type I and Type II drift.The trained Granite 3.3 2B recovers a large fraction of the recovery achieved by much larger GPT-4o, as stated in the supplied passage.
  • Future directions: Future work includes inverse-API correction for reversible writes, environment-grounded rewards, and evaluation on Type III drift and broader drift etiologies.These directions aim to extend diagnostic recovery toward a fully closed-loop, environment-aware safety layer.

CRediT Author Statement

The authors contributed across conceptualization, methodology, software, analysis, validation, investigation, writing, supervision, resources, and funding acquisition.

  • Ismail El Hamraoui contributed to conceptualization, methodology, software, formal analysis, and writing the original draft.
  • Sagar Jose contributed to conceptualization, methodology, software, validation, formal analysis, investigation, supervision, and writing-review and editing.
  • Nicolas Bureau contributed resources, funding acquisition, writing, and review and editing.
  • Robert Plana contributed conceptualization, validation, investigation, resources, and writing-review and editing.

Declaration of competing interest

The authors disclose their affiliations with Assystem entities and report no other commercial or financial relationships constituting a potential conflict of interest.

  • Four authors are affiliated with Assystem entities, and all authors declare no other commercial or financial relationships that could constitute a potential conflict of interest.Ismail El Hamraoui, Sagar Jose, and Nicolas Bureau work at the Digital Excellence Center, Assystem EOS; Robert Plana is Assystem’s Chief Technical Officer.

A Prompt Templates · B Judge Rubrics

The appendix defines five role-specific prompt templates and judge rubrics for structured drift diagnosis and recovery. Shared XML/JSON envelopes enforce format, while rubrics assess role-specific semantic correctness.

  • A Prompt Templates: All five node templates share one <reasoning> block followed by one <answer> JSON block, while NODE_ID, contextual fields, and expected keys vary.This shared envelope separates common structural requirements from node-specific inputs and outputs.
  • A.1 Classify Drift Prompt (n1): n1 classifies whether a reviewed step aligns with the task using the task, suspected onset, cause, and step content.Its output requires concise reasoning plus step, is_aligned, and why fields.
  • A.2 Detect Drift Operations Prompt: n2 identifies executed WRITE operations and task-misaligned READ operations, excluding attempted but failed operations.It receives n1’s verdict and reason and returns separate write_operations and read_operations fields.
  • A.3 Search API Prompt (n3): n3 deduplicates applications involved in n2’s writes so their documentation can support reversibility assessment.The output is an apps_name list of applications derived from the detected write operations.
  • A.4 Evaluate Risk Prompt (n4): n4 partitions detected writes into reversible and non-reversible categories using API documentation and inaccessible-data constraints.Its structured output contains write_reversible and write_not_reversible fields.
  • A.5 Aggregate and Decision Prompt: n5 chooses rollback_before_drift or escalate_human from prior evidence, targeting the step before the earliest confirmed drift when rollback is appropriate.Both outcomes require a short factual report, with rollback_to_step and report included in the structured arguments.
  • B Judge Rubrics: Each node’s judge rubric evaluates role-specific content quality alongside structural rewards, requiring grounded reasoning, consistent fields, evidence-based classifications, and actionable decisions.Rubrics also exclude failed or invented operations, prevent duplicate applications and overlapping risk partitions, and constrain rollback targets.

C Training Hyperparameters

The training configuration is specified in Table 5 and is shared across Granite 3.3 2B and Qwen 2.5 1.5B, with only their base checkpoints differing.

  • Table 5 provides the full training configuration for the approach.
  • The same configuration is used for Granite 3.3 2B and Qwen 2.5 1.5B.
  • The models differ only in their base checkpoints.
Loading 2608.14109v1…