Inference-Time Policy Alignment for Fair Reinforcement Learning

2026-07-31

AI Summary

This work addresses the challenge that existing reinforcement learning agents struggle to adapt to unknown fairness preferences after deployment, typically requiring retraining. To overcome this limitation, the study introduces— for the first time—the concept of inference-time alignment from large language models into reinforcement learning, proposing an inference-time policy alignment framework that requires no modification to pre-trained policy parameters. The method dynamically reshapes the action probability distribution through multiplicative policy shaping guided by action-specific welfare scores, thereby achieving desired fairness objectives. Experimental results demonstrate that the proposed framework significantly improves welfare-oriented fairness metrics across multiple tasks while incurring negligible degradation in original task performance, and it remains compatible with any deep reinforcement learning agent.

Abstract

Deep reinforcement learning (RL) agents achieve strong performance by optimizing scalar reward functions. However, once deployed, the policies of these RL agents are often rigid and costly to adapt to new performance criteria. For instance, an agent trained to maximize expected cumulative reward may not accommodate previously unknown stakeholder preferences. Existing approaches to achieve fairness, a type of preference, in RL typically assume that such preferences are known a priori and require complete retraining of the policy under a fairness-oriented metric. Inspired by inference-time alignment in large language models, we investigate the problem of steering a pretrained RL policy toward welfare-based fairness objectives at inference time without updating the base policy's parameters. We formalize inference-time fairness alignment as a policy shaping problem and propose a multiplicative policy shaping framework that adjusts action probabilities using action-dependent welfare scores, thus requiring no modification to the base policy. Our framework is general and compatible with any deep RL agent. Through extensive experiments across multiple domains, we demonstrate that inference-time policy shaping substantially improves welfare-based fairness objectives while preserving core task performance.

Problem

• fairness
• reinforcement learning
• inference-time alignment
• policy adaptation
• welfare-based objectives

Innovation

• inference-time alignment
• fair reinforcement learning
• policy shaping
• welfare-based fairness
• multiplicative adjustment