Robust General Utility for Reinforcement Learning

2026-08-04

AI Summary

This work addresses the lack of robustness in reinforcement learning (RL) arising from mismatches between utility functions used during deployment and training. To tackle this issue, the paper introduces a minimax universal utility RL framework that, for the first time, incorporates robustness into the general utility setting. The approach models utility misspecification via a utility uncertainty set, unifying paradigms such as reward-robust RL and constrained RL under a single formulation. For both concave and non-concave utility settings, the authors develop tailored algorithms—projected stochastic gradient descent-ascent and stochastic proximal hypergradient methods, respectively—ensuring convergence guarantees. Empirical evaluations on large language model safety alignment and exploration-maximization tasks validate the theoretical predictions, demonstrating both convergence and effectiveness of the proposed framework.

Abstract

Reinforcement learning (RL) with general utility extends classic RL by optimizing an arbitrary utility functional of the policy-induced occupancy measure, thereby enabling a broader range of applications. However, previous work on general utility RL typically assumes the evaluation utility is fixed and correctly specified. In practice, the utility used at deployment can deviate from the training one, creating a robustness gap that prior work does not address. Motivated by this, we propose robust general-utility RL, a minimax learning framework that trains policies against utility misspecification within a prescribed uncertainty set. Our framework strictly generalizes standard general-utility RL while also providing a unified view of many existing RL frameworks, including reward-robust RL and constrained RL, through appropriate choices of the utility uncertainty set. We further develop provably convergent stochastic algorithms for two regimes. For concave utilities, we develop a projected stochastic gradient descent-ascent method and establish stationarity guarantees. For the more challenging nonconcave regime, we propose a stochastic prox-extragradient algorithm that mitigates ill-posed behavior induced by nonconcavity, with convergence guarantees to approximate first-order stationarity. Experiments on LLM safety alignment and exploration maximization tasks further corroborate the convergence behavior consistent with our theory.

Problem

• robustness
• general utility
• reinforcement learning
• utility misspecification
• uncertainty set

Innovation

• robust reinforcement learning
• general utility
• minimax optimization
• nonconcave optimization
• stochastic prox-extragradient