The Trajectory Alignment Coefficient in Two Acts: From Reward Tuning to Reward Learning

TL;DR

This paper introduces the Trajectory Alignment Coefficient (TAC) as a tool to improve reward function design in reinforcement learning, supporting practitioners and enabling reward learning through a differentiable approximation called Soft-TAC.

Contribution

It presents TAC as a metric for aligning reward preferences with expert judgments and develops Soft-TAC for training reward models directly from human preferences.

Findings

TAC helps practitioners create more effective reward functions with lower workload.

Reward models trained with Soft-TAC better capture preference-specific objectives.

Using TAC as a reward learning objective yields policies with more distinct behaviors.

Abstract

The success of reinforcement learning (RL) is fundamentally tied to having a reward function that accurately reflects the task objective. Yet, designing reward functions is notoriously time-consuming and prone to misspecification. To address this issue, our first goal is to understand how to support RL practitioners in specifying appropriate weights for a reward function. We leverage the Trajectory Alignment Coefficient (TAC), a metric that evaluates how closely a reward function's induced preferences match those of a domain expert. To evaluate whether TAC provides effective support in practice, we conducted a human-subject study in which RL practitioners tuned reward weights for Lunar Lander. We found that providing TAC during reward tuning led participants to produce more performant reward functions and report lower cognitive workload relative to standard tuning without TAC. However, the study also underscored that manual reward design, even with TAC, remains labor-intensive. This limitation motivated our second goal: to learn a reward model that maximizes TAC directly. Specifically, we propose Soft-TAC, a differentiable approximation of TAC that can be used as a loss function to train reward models from human preference data. Validated in the racing simulator Gran Turismo 7, reward models trained using Soft-TAC successfully captured preference-specific objectives, resulting in policies with qualitatively more distinct behaviors than models trained with standard Cross-Entropy loss. This work demonstrates that TAC can serve as both a practical tool for guiding reward tuning and a reward learning objective in complex domains.