Evaluation-Aware Reinforcement Learning

TL;DR

This paper addresses the evaluation bottleneck in reinforcement learning, where high variance or bias in policy evaluation hampers deployment of safety- and performance-critical policies. It introduces Evaluation-Aware Reinforcement Learning (EvA-RL), which jointly optimizes deployment performance and evaluability by leveraging an assessment environment and a value-prediction penalty $\beta$ in the objective $\max_{\pi} \mathbb{E}_{s \sim \mu_D}[V_D^{\pi}(s) - \beta (V_D^{\pi}(s) - \hat{V}_D^{\pi}(s))^2]$. To mitigate the trade-off between returns and evaluation accuracy, EvA-RL co-learns a transformer-based value predictor alongside the policy, enabling the predictor to carry part of the evaluation burden. Theoretical analysis shows a convex relaxation yields an upper bound on the Bellman-consistent return and reveals a monotone trade-off with fixed predictors; empirically, EvA-RL substantially reduces evaluation error across discrete and continuous domains while maintaining competitive returns compared to standard RL and OPE baselines. This work lays a foundation for evaluating RL policies during training and points toward evaluation-aware methods as a principled direction for safer, more efficient RL development.

Abstract

Policy evaluation is often a prerequisite for deploying safety- and performance-critical systems. Existing evaluation approaches frequently suffer from high variance due to limited data and long-horizon tasks, or high bias due to unequal support or inaccurate environmental models. We posit that these challenges arise, in part, from the standard reinforcement learning (RL) paradigm of policy learning without explicit consideration of evaluation. As an alternative, we propose evaluation-aware reinforcement learning (EvA-RL), in which a policy is trained to maximize expected return while simultaneously minimizing expected evaluation error under a given value prediction scheme -- in other words, being "easy" to evaluate. We formalize a framework for EvA-RL and design an instantiation that enables accurate policy evaluation, conditioned on a small number of rollouts in an assessment environment that can be different than the deployment environment. However, our theoretical analysis and empirical results show that there is often a tradeoff between evaluation accuracy and policy performance when using a fixed value-prediction scheme within EvA-RL. To mitigate this tradeoff, we extend our approach to co-learn an assessment-conditioned state-value predictor alongside the policy. Empirical results across diverse discrete and continuous action domains demonstrate that EvA-RL can substantially reduce evaluation error while maintaining competitive returns. This work lays the foundation for a broad new class of RL methods that treat reliable evaluation as a first-class principle during training.

Figures (11)

• Figure 1: Value Prediction Transformer: A transformer encoder that takes as input (i) the start-states of the assessment environment -- $\{s^1, s^2, \ldots, s^k\}$, (ii) the returns corresponding to the rollouts of policy $\pi$ starting at these states -- $\{g(h_1), g(h_2), \ldots, g(h_k)\}$, and (iii) a query state $s \in \mathcal{S}_D$ and outputs an estimate for the value of the query state -- $\hat{V}_D^\pi(s)$. The JAX implementation is provided in appendix \ref{['sec:appendix_pred_transformer']}.
• Figure 1: EvA-RL often produces more accurate value estimates than OPE estimators in continuous-action environments.
• Figure 2: EvA-RL with a pre-trained value predictor. Top: normalized returns compared to the standard RL baseline $(\beta = 0)$. Bottom: predictor MAE. With a frozen predictor, increasing $\beta$ decreases MAE but also reduces returns. Co-learning allows for better performance while maintaining low evaluation error.
• Figure 3: Co-learned value predictor vs. OPE estimators for EvA-RL policies. The predictor consistently achieves lower evaluation error showcasing the usefulness of the behavior-conditioned value prediction.
• Figure 4: Comparison of standard RL (A2C + OPE) and EvA-RL in discrete-action environments. Top: normalized discounted returns (relative to standard RL; higher is better). Bottom: mean absolute error (MAE) of state-value estimates (lower is better). EvA-RL nearly matches standard RL in its performance while substantially reducing the evaluation error.

Theorems & Definitions (25)

• Definition 1
• Lemma 3.1
• proof
• Theorem 3.1
• proof
• Corollary 3.1
• proof
• Proposition 3.1
• proof
• Lemma 3.2