[Paper Review] Non-Markovian Reward Modelling from Trajectory Labels via Interpretable Multiple Instance Learning

Joseph Early, Tom Bewley+2 more|arXiv (Cornell University)|May 30, 2022

Time Series Analysis and Forecasting4 citations

TL;DR

This paper proposes a novel non-Markovian reward modelling approach using interpretable multiple instance learning (MIL), where trajectories are treated as bags and state-action pairs as instances to model temporal dependencies in human feedback. The method uses LSTM-based MIL models to reconstruct hidden state dynamics and reward functions from return-labeled trajectories, achieving high accuracy and enabling RL agents to solve non-Markovian tasks with performance matching or exceeding access to ground-truth rewards.

ABSTRACT

We generalise the problem of reward modelling (RM) for reinforcement learning (RL) to handle non-Markovian rewards. Existing work assumes that human evaluators observe each step in a trajectory independently when providing feedback on agent behaviour. In this work, we remove this assumption, extending RM to capture temporal dependencies in human assessment of trajectories. We show how RM can be approached as a multiple instance learning (MIL) problem, where trajectories are treated as bags with return labels, and steps within the trajectories are instances with unseen reward labels. We go on to develop new MIL models that are able to capture the time dependencies in labelled trajectories. We demonstrate on a range of RL tasks that our novel MIL models can reconstruct reward functions to a high level of accuracy, and can be used to train high-performing agent policies.

Motivation & Objective

• To generalize reward modelling beyond the Markovian assumption by accounting for temporal dependencies in human trajectory evaluations.
• To address the psychological and practical limitations of independent, instantaneous reward assessments in human feedback.
• To develop a method that reconstructs both hidden state dynamics and reward functions from trajectory return labels.
• To enable RL agents to learn high-performing policies using only return-labeled trajectories, without access to ground-truth rewards or hidden states.
• To provide interpretability tools for verifying and understanding the learned reward functions and hidden state dynamics.

Proposed method

• Formalizes non-Markovian reward modelling as a multiple instance learning (MIL) problem, where trajectories are bags and state-action pairs are instances.
• Uses long short-term memory (LSTM) networks to model hidden state dynamics across trajectories and infer instance-level reward predictions.
• Designs novel MIL model architectures that jointly learn reward functions and hidden state representations from bag-level return labels.
• Applies interpretability techniques to analyze and verify the internal representations and decision-making processes of the learned models.
• Trains models end-to-end on trajectory return labels, enabling inference of both reward and hidden state trajectories without supervision on instance-level labels.
• Integrates the learned reward models into RL training pipelines to evaluate policy performance on non-Markovian tasks.

Experimental results

Research questions

• RQ1 Can multiple instance learning effectively model non-Markovian reward functions that depend on hidden state dynamics not captured by visible states?
• RQ2 How well can MIL-based models reconstruct ground-truth hidden states and reward functions from return-labeled trajectories?
• RQ3 Can RL agents trained with the predicted rewards achieve performance comparable to or better than those with access to true rewards and hidden states?
• RQ4 How robust are the proposed MIL models to label noise and distributional shifts in trajectory data?
• RQ5 What insights can interpretability analysis provide into the internal representations and decision-making of the learned reward models?

Key findings

• The proposed MIL-based reward models accurately reconstruct ground-truth reward functions and hidden state trajectories in synthetic non-Markovian environments.
• The models achieve high return prediction accuracy, with performance comparable to or exceeding that of models trained with access to true hidden states.
• RL agents trained using the predicted rewards from the MIL models solve non-Markovian tasks with performance matching or exceeding that of agents with access to true rewards.
• Interpretability analysis reveals that the models learn meaningful hidden state dynamics, such as key collection and treasure region entry, aligning with human-inferred trajectory semantics.
• The models demonstrate robustness to label noise, maintaining strong performance even when return labels are corrupted or inconsistent.
• In the Lunar Lander task, the model learned practical improvements over the oracle reward function, such as rewarding safe hover positions away from the edge, which contributed to higher actual performance on the hover task.