PA3: Policy-Aware Agent Alignment through Chain-of-Thought

Shubhashis Roy Dipta; Daniel Bis; Kun Zhou; Lichao Wang; Benjamin Z. Yao; Chenlei Guo; Ruhi Sarikaya

Home/Archive/arXiv

arXiv · 2603.14602

PA3: Policy-Aware Agent Alignment through Chain-of-Thought

Shubhashis Roy Dipta·Daniel Bis·Kun Zhou·Lichao Wang·Benjamin Z. Yao·Chenlei Guo·Ruhi Sarikaya

View original source ↗Read source document ↗Export metadata ↓

Abstract

Conversational assistants powered by large language models (LLMs) excel at tool-use tasks but struggle with adhering to complex, business-specific rules. While models can reason over business rules provided in context, including all policies for every query introduces high latency and wastes compute. Furthermore, these lengthy prompts lead to long contexts, harming overall performance due to the 'needle-in-a-haystack' problem. To address these challenges, we propose a multi-stage alignment method that teaches models to recall and apply relevant business policies during chain-of-thought reasoning at inference time, without including the full business policy in-context. Furthermore, we introduce a novel Policy Recall reward based on the Jaccard score and a Hallucination Penalty for GRPO training. Altogether, our best model outperforms the baseline by 16 points and surpasses comparable in-context baselines of similar model size by 3 points, while using 40% fewer words.

Explore related subjects

cs.CL cs.AI cs.LG

Keep this discovery

BibTeX RIS

Shubhashis Roy Dipta, Daniel Bis, Kun Zhou, Lichao Wang, Benjamin Z. Yao, Chenlei Guo, Ruhi Sarikaya. 2026-08-31. PA3: Policy-Aware Agent Alignment through Chain-of-Thought. https://arxiv.org/abs/2603.14602

Cite the original work for its findings. Save a collection to share your selection of sources.

Discover connections

cs.AI subject cs.CL subject cs.LG subject

Connections use source metadata and explicit phrase matches, not verified experimental comparisons.

KEEP EXPLORING

Related discoveries

arXiv2026

Language-Guided Tuning: Configuration Optimization for Automated ML Research

Configuration optimization remains a critical bottleneck in machine learning, requiring coordinated tuning across model architecture, training strategy, feature engineering, and hyperparameters. Traditional approaches treat these dimensions independently and lack interpretability, while recent automated methods struggle with dynamic adaptability and semantic reasoning about optimization decisions. We introduce Language-Guided Tuning (LGT), a framework that employs multi-agent Large Language Models to automatically optimize configurations through natural language reasoning. We apply textual feedback signals that complement numerical optimization by providing semantic understanding of training dynamics and configuration interdependencies. LGT coordinates three specialized agents: an Advisor that proposes configuration changes, an Evaluator that assesses progress, and an Optimizer that refines the decision-making process, creating a self-improving feedback loop. Through comprehensive evaluation on seven diverse datasets, LGT demonstrates substantial improvements over traditional optimization methods while maintaining high interpretability.

cs.AI↗

arXiv2026

G-Loss: Graph-Guided Fine-Tuning of Language Models

Traditional loss functions, including cross-entropy, contrastive, triplet, and su pervised contrastive losses, used for fine-tuning pre-trained language models such as BERT, operate only within local neighborhoods and fail to account for the global semantic structure. We present G-Loss, a graph-guided loss function that incorporates semi-supervised label propagation to use structural relationships within the embedding manifold. G-Loss builds a document-similarity graph that captures global semantic relationships, thereby guiding the model to learn more discriminative and robust embeddings. We evaluate G-Loss on five benchmark datasets covering key downstream classification tasks: MR (sentiment analysis), R8 and R52 (topic categorization), Ohsumed (medical document classification), and 20NG (news categorization). In the majority of experimental setups, G-Loss converges faster and produces semantically coherent embedding spaces, resulting in higher classification accuracy than models fine-tuned with traditional loss functions.

cs.CL↗

arXiv2026

Stay Within Your Bounds: Distance-Guided Decoding for Guaranteed Context-Free Grammar Compliance

Grammar-constrained decoding helps large language models produce syntactically valid structured outputs, such as code, JSON, and SQL. For context-free grammars, many practical decoders enforce local prefix feasibility: each token must keep the current prefix extendable to some valid completion. Yet, under tokenizer-grammar mismatch and finite token budgets, feasible prefixes may still fail to reach acceptance. We propose a lookahead-guided decoding framework for context-free grammars based on pushdown automata. Offline, we compute bounded pushdown summaries with reachability labels and upper-bound distances to acceptance. Online, these estimates guide horizon-aware pruning and beam search. The resulting decoder is syntactically sound: every output is accepted by the target grammar. Experiments on JSON, SQL, and Linear Temporal Logic (LTL) show both consistent syntactic validity and improved completion quality over existing baselines.

cs.AI↗