Aligning Machiavellian Agents: Behavior Steering via Test-Time Policy Shaping

TL;DR

This work tackles the misalignment risk of reward-optimizing AI agents by introducing test-time policy shaping, a lightweight, model-driven method that uses attribute classifiers to steer actions without retraining. It trains per-attribute classifiers on the MACHIAVELLI benchmark and interpolates their outputs with the base RL policy to control ethical attributes at inference. The approach demonstrates substantial reductions in ethical violations and power-seeking across diverse text-based games, often matching or outperforming training-time alignment baselines while revealing the trade-offs between reward and alignment. It also provides tools to analyze attribute correlations and even reverse prior alignment, highlighting the method's flexibility and practical impact for scalable, context-sensitive AI alignment.

Abstract

The deployment of decision-making AI agents presents a critical challenge in maintaining alignment with human values or guidelines while operating in complex, dynamic environments. Agents trained solely to achieve their objectives may adopt harmful behavior, exposing a key trade-off between maximizing the reward function and maintaining alignment. For pre-trained agents, ensuring alignment is particularly challenging, as retraining can be a costly and slow process. This is further complicated by the diverse and potentially conflicting attributes representing the ethical values for alignment. To address these challenges, we propose a test-time alignment technique based on model-guided policy shaping. Our method allows precise control over individual behavioral attributes, generalizes across diverse reinforcement learning (RL) environments, and facilitates a principled trade-off between ethical alignment and reward maximization without requiring agent retraining. We evaluate our approach using the MACHIAVELLI benchmark, which comprises 134 text-based game environments and thousands of annotated scenarios involving ethical decisions. The RL agents are first trained to maximize the reward in their respective games. At test time, we apply policy shaping via scenario-action attribute classifiers to ensure decision alignment with ethical attributes. We compare our approach against prior training-time methods and general-purpose agents, as well as study several types of ethical violations and power-seeking behavior. Our results demonstrate that test-time policy shaping provides an effective and scalable solution for mitigating unethical behavior across diverse environments and alignment attributes.

Figures (13)

• Figure 1: Overview of our proposed alignment approach using test-time policy shaping. Given a scenario, ethical attribute classifiers predict the likelihood of different attributes for each available action. These predictions are then used to adjust an agent’s policy during inference to discourage actions misaligned with ethical target attributes, e.g. avoiding killing.
• Figure 2: Distribution of ethical attributes in the MACHIAVELLI benchmark across the 10 chosen test games. See the Appendix for more details about the selection process.
• Figure 4: Alignment results of the RL, Oracle, and policy shaping RL-$\alpha 0.2$, RL-$\alpha 0.8$, and RL-$\alpha 1.0$ agents per the top five ethical violations and power. Oracle and RL-$\alpha$ agents are steered to minimize deception (denoted as "dec."), resulting in a decrease of deception as $\alpha$ increases. The RL-$\alpha 1.0$ agent achieves the best score, closest to the Oracle.
• Figure 5: Pareto front showing the trade-off of points (i.e., reward) and violation score of RL agents with our policy-shaping approach applied per top-5 ethical violation.
• Figure 6: Correlation between ethical attributes when applying policy shaping. The bottom half of the matrix illustrates the results of agents minimizing attributes, and the top half illustrates maximizing attributes. Attribute names are abbreviated, with power-seeking attributes denoted by "p.", "non.p" is non-physical harm, and "int.h" is intending harm.