The Trajectory Alignment Coefficient in Two Acts: From Reward Tuning to Reward Learning
Calarina Muslimani, Yunshu Du, Kenta Kawamoto, Kaushik Subramanian, Peter Stone, Peter Wurman
TL;DR
This work investigates the Trajectory Alignment Coefficient (TAC) as a practical aid for reward design in reinforcement learning and as an objective for reward learning. It demonstrates that providing TAC feedback improves final policy performance and reduces cognitive workload during manual reward tuning in Lunar Lander, while also introducing Soft-TAC, a differentiable TAC surrogate used to learn rewards directly from human preferences. Soft-TAC is validated in Lunar Lander and Gran Turismo 7, where learned reward models better capture human objectives and induce qualitatively distinct behaviors (e.g., fast vs. aggressive vs. timid driving). The results indicate that TAC can both guide reward design and serve as a robust learning objective, with broader implications for human-aligned RL in complex tasks, though the approach is currently demonstrated with linear reward functions. Future directions include extending TAC to black-box reward models and exploring informative feature identification to further reduce design effort and improve generalization.
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.