Non-Markovian Reward Modelling from Trajectory Labels via Interpretable Multiple Instance Learning
Anonymous Author(s)
AffiliationAddress email
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 onagent behaviour.
In this work, we remove this assumption, extending RM to capturetemporal dependencies in human assessment of trajectories.
We show how RM canbe approached as a multiple instance learning (MIL) problem, where trajectoriesare treated as bags with return labels and steps within the trajectories are instanceswith unseen reward labels.
We go on to develop new MIL models that are able tocapture the time dependencies in labelled trajectories.
We demonstrate on a rangeof RL tasks that our novel MIL models can reconstruct reward functions to a highlevel of accuracy and can be used to train high-performing agent policies
Introduction
There is growing consensus around the view that aligned and beneficial AI requires a reframingof objectives as being contingent, uncertain, and learnable via interaction with humans [35].
Inreinforcement learning (RL), this proposal has found one formalisation in reward modelling (RM):the inference of agent objectives from human preference information such as demonstrations, pairwisechoices, approval labels, and corrections [29].
Prior work in RM typically assumes that a humanevaluates the return (quality) of a sequential trajectory of agent behaviour by summing equal andindependent reward assessments of instantaneous states and actions, with the aim of RM beingto reconstruct the underlying reward function.
However, in reality the human’s experience of atrajectory is likely to be temporally-extended (e.g., via a video clip [12] or real-time observation),which opens the door to dependencies between earlier events and the assessment of later ones.
Theindependence assumption may be both psychologically unrealistic given human memory limitations[26], and technically naïve given the difficulty of building complete instantaneous state representations[25].
We thus seek to generalise RM to allow for temporal dependencies in human evaluation, bypostulating hidden state information that accumulates over a trajectory.
Reconstruction of the human’spreferences now requires the modelling of hidden state dynamics alongside the reward function itself.
In tackling this generalised problem, we identify a structural isomorphism between RM (specificallyfrom trajectory return labels) and the established field of multiple instance learning (MIL) [10].
Trajectories are recast as bags and constituent state-action pairs as instances , which collectivelycontribute to labels provided at the bag level by interacting in potentially complex ways.
Thismapping inspires a range of novel MIL model architectures that use long short-term memory (LSTM)modules [19] to recover the hidden state dynamics, and learn instance-level reward predictions fromreturn-labelled trajectories of arbitrary length.
In experiments with synthetic oracle labels, we showthat our novel MIL RM models can accurately reconstruct ground-truth hidden states and rewardfunctions for non-Markovian tasks, and can be straightforwardly integrated into RL agent training to achieve performance matching, or even exceeding, that of agents with direct access to true hiddenstates and rewards.
Our contributions are as follows:
We generalise RM to handle non-Markovian rewards that depend on hidden features of the
environment or the psychology of the human evaluator in addition to visible states/actions.2.
We identify a structural connection between RM and MIL, creating the opportunity totransfer concepts and methods between the two fields.
We propose novel LSTM-based MIL models for this generalised RM problem, and developinterpretability techniques for understanding and verifying the learnt reward functions.
We compare our proposed models to existing MIL baselines on five non-Markovian tasks,evaluating return prediction, reward prediction, robustness to label noise
and interpretability.
We demonstrate that the hidden state and reward predictions of our MIL RM models can beused by RL agents to solve non-Markovian tasks.
The remainder of this work is as follows.
Section 2 discusses related work in RM and MIL, Section3 gives a formal problem definition and describes our MIL-inspired methodology, and Section 4presents our experiments and results.
We discuss key findings in Section 5, and Section 6 concludes.
Background and Related Work
Reward Modelling RM [29] aims to infer a reward function from revealed human preferenceinformation such as demonstrations [32], pairwise choices [12], corrections [4], good/bad/neutrallabels [33], or combinations thereof [24].
Most prior work assumes a human evaluates a trajectory bysumming independent rewards for each state-action pair, but in practice their experience is likely tobe temporally-extended (e.g., via a video clip), creating the opportunity for dependencies to emergebetween earlier events and the assessment of later ones.
As noted by Chan et al.
[11] and Bewleyand Lecue [8], such dependencies may arise from cognitive biases such as anchoring, prospect bias,and the peak-end rule [26], but they could equally reflect rational drivers of human preferences notcaptured by the state representation.
Some efforts have been made to model temporal dependencies,such as a discrete psychological mode which evolves over consecutive queries about hypotheticaltrajectories [6], or a monotonic bias towards more recently-viewed timesteps due to human memorylimitations [28].
Elsewhere, Shah et al.
[36] use human demonstrations and binary approval labelsto learn temporally-extended task specifications in logical form.
In comparison to these restrictedexamples, our work provides a more general approach to capturing temporal dependencies in RM.
Non-Markovian Rewards In the canonical RL problem setup of a Markov decision process (MDP),rewards depend only on the most recent state-action pair.
In a non-Markovian reward decision process(NMRDP) [2], rewards depend on the full preceding trajectory [2].
NMRDPs can be expanded into
MDPs (and thus solved by RL) by augmenting the state with a hidden state that captures all rewardrelevant historical information, but this is typically not known a priori .
Data-driven approaches tolearning NMRDP expansions [21] often make use of domain-specific propositions and temporal logicoperators [3, 39, 41].
Outside of the RM context, recurrent architectures such as LSTMs have beenused in NMRDPs to reduce reliance on pre-specified propositions [23].
They also have a long historyof use in partially observable MDPs, where dynamics are also non-Markovian [5, 46, 17].
Multiple Instance Learning In MIL [10], datasets are structured as collections of bags X i ∈ X ,each of which is comprised of instances { x i 1 , . . .
, x ik } and has an associated bag-level label Y i andinstance-level labels { y i 1 , . . .
, y ik } .
The aim is to construct a model that learns solely from bag labels;instance labels are not available during training, but may be used later to evaluate instance-levelpredictions.
The simplest MIL approaches assume that instances are independent and that the bag isunordered, but models exist for capturing various types of instance dependencies [45, 22, 42].
LSTMshave emerged as a natural architecture for modelling temporal dependencies among ordered bags,where they can be utilised to aggregate instance information into an overall bag representation.
Theyhave previously been applied to standard MIL benchmarks [44], as well as specific problems such as
Chinese painting image classification [30].
As we discuss in Section 3.2, these existing models aresomewhat unsuitable for use in RM, leading us to propose our own novel model architectures.
Methodology
In this section, we present the core methodology of our work.
We formally define the new paradigmof non-Markovian RM (Section 3.1), before drawing on the MIL literature to propose models thatcan be used to solve this generalised problem (Section 3.2).
We then go on to discuss how we can useour learnt RM models for training RL agents on non-Markovian tasks (Section 3.3).
3.1 Formal Definition of Non-Markovian RM
Consider an agent interacting with an environment with Markovian dynamics.
At discrete time t ,the current environment state s t ∈ S and agent action a t ∈
A condition the next environment states t +1 according to the dynamics function D : S ×
A → ∆ S .
A trajectory ξ ∈ Ξ is a sequence ofstate-action pairs, ξ = (( s 0 , a 0 ) , ..., ( s T − 1 , a T − 1 )) , and a human’s preferences about agent behaviourrespect a real-valued return function G : Ξ → R .
In traditional (Markovian) RM, return is assumedto decompose into a sum of independent rewards over state-action pairs, G ( ξ )
= (cid:80) T − 1 t =
R ( s t , a t ) ,and the aim is to reconstruct R ′
≈ R from possibly-noisy sources of preference information.
In ourgeneralised non-Markovian model, we consider the human to observe a trajectory sequentially andallow for the possibility of hidden state information that accumulates over time and parameterises R :
G ( ξ ) =
T − 1 (cid:88) t =
R ( s t , a t , h t +1 ) where h t +1 =
δ ( h t , s t , a t ) , (1) δ is a hidden state dynamics function, and h 0 is a fixed value for the initial hidden state.
Reconstructionof the human’s preferences now requires the estimation of δ ′ ≈ δ and h ′
≈ h 0 alongside R ′
≈ R .
We can interpret the hidden information as an external feature of the environment that is visible tothe human but excluded from the state, or as a psychological state of the person themselves, throughwhich their response to a new observation is influenced by what they have seen already.
Appendix Belaborates on this discussion, presenting motivating use cases and limitations of non-Markovian RM.
We visualise the difference between Markovian and non-Markovian RM in Figure 1.
In this work, we focus on the simplest and most explicit form of preference information: directlabelling of returns G ( ξ i ) for a dataset of N trajectories { ξ i } Ni =1 .
We aim to solve the reconstructionproblem by minimising the squared error in predicted returns:
[Equation]
where h
[Equation]
i
[Equation]
We observe that Equation 2 perfectly matches the definition of a MIL problem.
Each trajectory ξ ican be considered as an ordered bag of instances (( s i, 0 , a i, 0 ) , ..., ( s i,T i − 1 , a i,T i − 1 )) with unobservedinstance labels R ( s i,t , a
i,t , h i,t +1 ) , an observed bag label G ( ξ i ) =
(cid:80) T − 1 t =
R ( s i,t , a i,t , h i,t +1 ) , andtemporal instance interactions via the changing hidden state h
i,t .
This correspondence motivates usto review the space of existing MIL models (specifically those that model temporal dependenciesamong instances) to provide a starting point
3.2 MIL RM Architectures
The MIL literature contains a variety of architectures for handling temporal instance dependencies,including graph neural networks (GNNs) [42] and transformers [37].
While effective for manyproblems, such architectures are an unnatural fit to non-Markovian RM as they contain no directanalogue of a hidden state h
′ carried forward in time, instead handling dependencies via some variantof message-passing between instances.
LSTM-based MIL architectures [44, 30] provide a morepromising starting point since they explicitly represent both h ′ (implemented as a continuous-valuedvector) and its temporal dynamics function δ ′ (a particular arrangement of gating functions).
Starting from an existing LSTM-based MIL architecture, we propose two successive extensions aswell as a naïve baseline that cannot handle temporal dependencies.
All four architectures include afeature extractor (FE) for mapping state-action pairs into feature vectors and a head network (HN) thatoutputs predictions.
These architectures are depicted in Figure 2.
Note we use the same nomenclatureas [10] and [45]: embedding space approaches produce an overall bag representation that is used forprediction, while instance space approaches produce predictions for each instance in the bag andthen aggregate those predictions to a final bag prediction.
Base Case: Embedding Space LSTM (ESL)
This architecture, proposed by Wang et al.
[44],processes all instances in a bag sequentially and uses the final LSTM hidden state as a bag embedding.
This is fed into the HN, which predicts the bag label g ′ (return in the RM context).
Although thismodel can account for temporal dependencies, it does not inherently produce instance predictions(rewards), which require some post hoc analysis to recover.
While methods exist for computinginstance importance values as a form of interpretability [13], these are not guaranteed to sum to thebag label as stipulated by the reward-return formulation.
We propose a new method: at time t , thepredicted reward r ′ t is calculated by feeding the LSTM hidden state at times t − 1 and t into the HNto obtain two partial bag labels/returns g ′ t − 1 and g ′ t , and computing the difference of the two, i.e.,r ′ t =
g ′ t − g ′ t − 1 .
Note that we define g ′
= 0 .
Extension 1: Instance Space LSTM (ISL)
The preceding post hoc computation is rather inelegantand often yields poor reward predictions (see Appendix D and Section 4), likely because rewards arenever computed or back-propagated through during learning.
This leads us to propose an improvedarchitecture, which is structurally similar but differs in how network outputs are mapped onto RMconcepts, and places reward predictions on the back-propagation path.
Given the LSTM hidden stateat time t , the output of the HN is taken to be the instantaneous reward r ′ t rather than the partial return.
Rewards are computed sequentially for all timesteps in a trajectory and summed to give the returnprediction g ′ .
We thereby obtain a model that both handles temporal dependencies and producesexplicitly-learnt reward predictions.
Extension 2: Concatenated Skip Connection (CSC) Instance Space LSTM (CISL)
In both of thepreceding architectures, the LSTM hidden state h ′ t is the sole input to the HN.
This requires h ′ t torepresent all reward-relevant information from both the true hidden state h t and the latest state-actionpair s t − 1 , a t − 1 to achieve good performance.
To lighten the load on the LSTM, we further extendthe ISL model with a skip connection [18, 20] which concatenates the FE output onto the hiddenstate before feeding it to the HN.
In principle, this should allow the hidden state to solely focuson representing temporal dependencies.
As well as improving RM performance compared to anequivalent model without skip connections, we find in Section 5.1 that this modification tends toyield more interpretable and disentangled hidden state representations.
Markovian Baseline: Instance Space Neural Network (ISN) To quantify the cost of ignoringtemporal dependencies, we also run experiments with a baseline architecture that feeds only the FEoutput for each state-action pair into the HN, yielding fully-independent reward predictions which aresummed to give the return prediction.
This independent predict-and-sum architecture has precedencein both MIL, where it is referred to as mi-Net by Wang et al.
[45], and in RM, where it embodies thede facto standard Markovian reward assumption [12].
3.3 Training Agents with Non-Markovian RM Models
In this work, as in RM more widely, we are not solely interested in learning reward functions torepresent human preferences, but also in the downstream application of rewards to train agents’action-selection policies.
After optimising our LSTM-based models on offline trajectory datasets, wedeploy them at the interface between conventional RL agents and their environments.
Going beyondprior work, where a learnt model is used to either generate a reward signal for an agent to maximise[12] or augment its observed state representation with hidden state information [21], our models servea dual role, providing both rewards and state augmentations.
Figure 3 describes this setup in detail.
following the relevant steps from Figure 2, and the reward is also sent to the agent.
In the language of NMRDPs, the hidden state augmentation expands the agent’s learning problem into an MDP by providing the additional information required to make the rewards Markovian.
Note that unlike during learning of the MIL models, return predictions are never required.
Experiments and Results
After initially validating our models on several toy datasets (see Appendix D), we focus the bulk ofour evaluation on five RL tasks.
As running experiments with people is costly, we use the standard
RM approach of generating synthetic preference data (here trajectory return labels) using ground-truthoracle reward functions [12] (for a discussion comparing the use of oracle and human labels, see
Appendix C).
Unlike prior work, these oracle reward functions depend on historical informationthat cannot be recovered from the instantaneous environmental state, thereby emulating the disparitybetween the information that a human evaluator possesses while viewing a trajectory sequentially,and that contained in the state alone.
In this section, we introduce our RL tasks (Section 4.1), evaluatethe quality of reward reconstruction (Section 4.2), investigate the use of MIL RM models for agenttraining (Section 4.3), and evaluate their robustness to label noise (Section 4.4).
4.1 RL Task Descriptions
We apply our methods to five non-Markovian RL tasks, the first four of which are within a common2D navigation environment and are specifically designed to capture different kinds of non-Markovianstructure.
Each environment has two spawn zones and an episode time limit of T = 100 ; see Figure4.
In each case, the environment state contains the x, y position of the agent only.
The tasks involvemoving into a treasure zone, contingent on some hidden information that cannot be derived from thecurrent x, y position, but is instead a function of the full preceding trajectory.
In the first two cases thehidden information varies with time only, but in the other two it depends on the agent’s past positions.
Timer For times t ≤ 50 the treasure gives a reward of − 1 for each timestep that the agent spendsinside it, before switching to +1 thereafter.
Since time is not included in the environment state,recovering the reward function by only observing the agent’s current position is impossible.
Moving The Timer task only captures a binary change, therefore we generalise it to be continuous.
In this case, the treasure zone oscillates left and right at a constant speed.
Again, this is not capturedin the environment state, but can be recovered if the length of the preceding trajectory is known.
Key Before reaching the treasure zone, the agent must first enter a second zone to collect a key;otherwise it receives 0 reward.
As the key’s status is not captured in the environment state, a temporaldependency exists between the agent’s past positions and the reward it obtains from the treasure.
Charger We generalise the Key task by replacing the key zone with a charging zone that buildsup the amount of reward the agent will receive when it reaches the treasure.
Now, the reward isdependent on not only if the agent visits a zone (binary), but how long it is there (continuous).
For the fifth and most complex task, we adapt Lunar Lander from OpenAI Gym [9], adding thecondition that the lander should take off again and stably hover after 50 timesteps on the landingpad.
This is analogous to the Charger task but with a larger state-action space and longer episodes( T = 500 ).
Further details on the tasks and MIL model hyperparameters are given in Appendix E.
[Figure]
An important design decision for the LSTM-based models is the size of the hidden state, as it affectsboth performance and interpretability.
For all the above tasks, we know a priori that it is possible tocapture the temporal dependencies in at most two dimensions, so we constrain our models to use 2Dhidden states.
This allows us to visualise and interpret the hidden representations in Section 5.1.
4.2 Reward Modelling Results
Below we discuss the performance of the reward reconstruction for the different MIL RM modelson our five RL tasks.
For each task, we generate initial trajectories to form our MIL RM datasets(see Appendix E).
Results from MIL models trained on these trajectories are given in Table 1.
Weobserve that the CISL is on average the best performing model for predicting both trajectory returnsand timestep rewards.
While the ESL model performs best at predicting return on the Key and Lunar
Lander tasks, it struggles on the reward metric
(due to the use of a proxy post hoc method).
Also notethat the ISN that serves as our Markovian RM baseline performs very poorly on return prediction,indicating that these tasks indeed cannot be learnt without modelling temporal dependencies.
4.3 RL Training Results
Following the method in Section 3.3, we then train Soft Actor-Critic [15] (Lunar Lander) and Deep
Q-Network [31] (all others) RL agents to optimise the rewards learnt by the LSTM-based models.
We evaluate agent performance in a post hoc manner by passing its trajectories to the relevant oracle.
This evaluation provides an end-to-end measure of both reward reconstruction and policy learning,
[Table]
and is standard in RM [12].
We baseline against agents trained with access to: a) the oracle rewardfunction and the oracle hidden states, and b) just the oracle reward function without hidden states(i.e, using only the environment states that are missing information).
In Figure 5, we observe thatthe CISL model enables the best RL agent performance, coming closest to the oracle.
Interestingly,for the Timer and Lunar Lander tasks, the CISL model actually outperforms the use of the oracle,suggesting that the learnt hidden states are easier to exploit for policy learning than the raw oraclestate (we investigate what these models have learnt in Section 5.1).
Note the poor performance ofagents trained without hidden state information, which aligns with expectations.
Appendix F containsfurther details on agent training, and Appendix G contains additional analysis for Lunar Lander.
[Figure]
4.4 Robustness to Mislabelling
In this work, the return labels are provided by oracles rather than real people.
When using humanevaluators, there is likely to be uncertainty in the labels, and it is important to evaluate modelrobustness against such noise [28].
We implement noise through label swapping [34]; this ensuresthe marginal label distribution remains the same and does not include out-of-distribution returns.
In
Figure 6, we show how both return and reward prediction decay with noise levels increasing from0 (no labels swapped) to 0.5 (half swapped).
The rate, smoothness, and consistency (across threerepeats) of this decay varies between tasks, with decays in return prediction generally being smoother.
We observe that the CISL model remains the strongest predictor of both return and reward in themajority of cases, indicating a general robustness and providing evidence that the model should stillbe effective with imperfect human labels.
On all metrics aside from Timer reward loss (where the mixof negative and positive rewards makes the effect of noise especially unpredictable), a noise level ofat least 0.3 is required for the CISL to perform as badly as the ISN baseline does with no noise at all.
Discussion
In this section, we seek to interpret our MIL RM models, analysing the distribution of learnt hiddenstates (Section 5.1) as well as their temporal dynamics over the course of a trajectory (Section 5.2).
Finally, in Section 5.3 we discuss the limitations of this work and potential areas for future work.
Charger
[Figure]
n
M
[Figure]
L o ss
[Figure]
r
[Figure]
5.1 Hidden State Analysis
The primary purpose of RM is to perform accurate reward reconstruction to facilitate agent training,but there is a secondary opportunity to improve understanding of human preferences through interpretability analysis of the learnt models.
We can directly visualise the 2D LSTM hidden states of ouroracle experiments, which enables a qualitative comparison of the various model architectures (see Figure 7).
Visualising the hidden states with respect to the temporal dependencies indicates that the CISL has learnt insightful hidden state representations.
Breaking down the CISL hidden embeddings:for the Timer task, time is represented along a curve, with a sparser representation around t = 50 (thecrossover point when the treasure becomes positive).
For the Moving task, time is similarly capturedalong with an additional notion of the change in treasure direction from right to left.
For the Keytask, the binary state of no key vs key is separated, with additional partitioning based on x position,denoting the two different start points of the agent.
We discuss the Charger task in Section 5.2.
[Figure]
E [Figure] b e dd i n g
S p a c e
L S T
M
I n s t a n c e
S p a c e
L S T
M
C S C
I n s t a n c e
S p a c e
L S T
M
[Figure]
5.2 Trajectory Probing
We further interpret our models by visualising the learnt reward with respect to the environmentstate, and by using hand-specified probe trajectories to verify that the learnt hidden state transitionsmimic the true transitions.
We present the above for the CISL model on the Charger task in Figure 8(Appendix H.2 contains similar figures for the Timer, Moving, and Key tasks).
The top row showsthat the model has correctly learnt the relationships between position, charge, and reward (rewardincreases in the treasure zone as charge increases).
From the probes, we can see how the charge levelcan be recovered from the hidden states.
We also note that the inflection point between under-chargingand over-charging is captured, i.e., this is where the optimal charge level lies, subject to some noisebased on where the agent starts in the spawn zones.
Furthermore, with the Challenging probe, weobserve that the learnt hidden states align with the agent moving in and out of the charging zone.
[Figure]
n v i r o n m e n t S t a t e
T r a j e c t
o r y
H i dd e n S t a t e
T r a j e c t
o r y
[Figure]
OutInCharging
[Figure]
to charge after maximum charge of 1 is reached.
“ x n " labels indicate the agent remains in a position for n timesteps.
As in Figure 7, we use the best performing model according to reward reconstruction.
5.3 Limitations and Future Work
Although we analyse the performance of our methods in the presence of noisy labels in Section 4.4, amajor area of future work is to apply our methods to human labelling (for a discussion of this, see
Appendix C).
Another area of future work involves more complex environments, for example the useof tasks with image observations, similar to the Atari environments in Open AI
Gym [9].
Furthermore,we perform RM from either an offline dataset or from only one RL training iteration; an iterativebootstrapping approach with multiple RL + RM training iterations could lead to improved RL results.
There are also limitations with our MIL RM approach for the Lunar Lander task; see Appendix E.5for details and suggestions for future work.
More generally, we hope that our identification of thelink between RM and MIL may inspire a bidirectional transfer of tools and techniques.
Conclusion
We posed the problem of non-Markovian RM, which removes an unrealistic assumption about howhumans evaluate temporally-extended agent behaviours.
After identifying an isomorphism between
RM and MIL, we proposed and evaluated novel MIL-inspired models that allow us to reconstructnon-Markovian reward functions, augment agent training, and interpret their learnt representations.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 5.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A](b)
Did you include complete proofs of all theoretical results?
[N/A]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See Appendix
A
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendices D, E, and F
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Appendix A
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix A
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
Our source code and models, which are included in the supplementary material and will be made available publicly upon acceptance.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Implementation and Resource Details
This work was implemented in Python 3.8 / 3.10 and the machine learning functionality used PyTorch.
All required libraries for our work are given in a requirements.txt file.
Our code is included in the
Supplementary Material submission and will be made publicly accessible once the work is published.
The majority of MIL model training was carried out on a remote GPU service using a Volta V
Enterprise Compute GPU with 16GB of VRAM, which utilised CUDA v11.0 to enable GPU support.
For the Lunar Lander task, training each MIL model took a maximum of eight hours.
For the othertasks, this was a maximum of two hours.
Trained models are included alongside the code.
Fixed seedswere used to ensure consistency of dataset splits between training and testing; these are included inthe scripts that are used to run the experiments.
All our datasets were generated from code; both thescripts to generate the data and also the derived datasets themselves are included alongside our modeltraining code.
Dataset generation, as well as all RL agent training, was conducted on a second remote
GPU service using a compute node with two Nvidia Pascal P100 cards.
Data generation took amaximum of three hours per dataset.
Agent training for the 2D navigation tasks was computationallylight, requiring 8-12 minutes per 400-episode run, although we completed ten repeat runs for eachpermutation of task and MIL model architecture (one per MIL training repeat).
800-episode RL runsfor Lunar Lander took approximately two hours each.
Further details on executing the scripts toreproduce our results can be found in the README.md file in our code submission.
B Use Cases for Non-Markovian Reward Modelling
The non-Markovian reward formulation applies to cases where rewards depend on hidden stateinformation h t in addition to environment states s t and actions a t , and this information is a functionof previous state-action pairs but not vice versa (i.e., there is no causal path from h t to s t + k for anyk > 1 ).
This crucial caveat distinguishes the formulation from the more general class of partiallyobservable MDPs and demarcates the set of domains to which it can be applied: those involving asecondary Markovian system that “spectates" on events in the environment without intervening.
Inthe RM context, this secondary system is a black box (making its internal state h t hidden) and explicitrewards are unavailable, being replaced by sparser and potentially noisy form of reward-dependentfeedback (trajectory return labels in our work).
Below we identify three classes of use case which fitthis technical specification and provide one concrete example for each:
Ambiguous Subtasks Cases involving an extended task with a sequential structure, where it is
hard to formally define the conditions for subtask completion, but RM is feasible because a human
“knows it when they see it".
Here, the hidden state to be learnt represents the current subtask and any
auxiliary information needed to determine its completion status.
• Concrete Example: Using judges’ scores to learn a performative display (e
g., gymnastics,
aerobatics) chaining several manoeuvres whose start and end conditions are difficult to
formalise a priori .
This could be considered as an extension of the single backflip task
studied in the foundational RM work by Christiano et al.
Dependencies on Subjective Affect Cases where a human’s reward function is dependent on
their affective (emotional) status, which in turn depends on their prior experiences.
Assuming this
information is not directly available from the observed environment state, it must be learnt from data.
• Concrete Example: Using periodic satisfaction ratings to train a personal assistant robotwhose owner’s mood, needs and preferences vary from day to day.
These variations may
influence the preferred driving style of a chauffeur service or choice of evening meal.
Irrationalities/Cognitive Biases Cases where one or more forms of bias colour a human’s post
hoc rating of an observed trajectory, even if their instantaneously-experienced reward is Markovian.
Psychological studies of how humans aggregate immediate rewards into retrospective evaluations of
the quality of an experience find that a straight summation assumption is unrealistic, with subjectsexhibiting high sensitivity to contrast effects and recency bias (collectively termed the peak-end rule)[26], and factoring in anticipated future states in addition to those actually observed [1].
Here, thehidden state captures an aggregate representation of the biases at play in a given human’s evaluation.
• Concrete Example: Using customer “star ratings" to improve a holiday planning agentwhose recommendations aim to account for an unknown mix of biases such as the peak-endrule.
The agent may use the learnt bias model to prioritise key moments in a holiday when managing the travel schedule and budget, in order to maximise future customers’ star ratings.
Finally, we note that it is a matter of taste as to whether hidden state information is framed assituated inside a human evaluator’s mind or in the environment but only visible to the human .
Itis not technically necessary to decide between these two framings, as the mathematical problemof non-Markovian RM is equivalent.
From an agent’s perspective, a human evaluator is part of anaugmented environment, even if they never intervene directly to influence the state.
C Comparing Human and Oracle Labelling
Oracle-based experiments are often to used to evaluate RM methods, since they enable scalablequantitative validation [14, 16, 33]. However, we identify three concrete differences between ouroracle preference labelling method and realistic human labelling: 1) preference form, 2) preferencesparsity, and 3) preference noise. Preference form is the different ways of providing labels; in thiswork we used return values which are highly informative and easy to learn from, but it has long beenunderstood that humans find it easier to give less direct feedback, such as pairwise rankings [12]or good/bad/neural labels [27].
Preference sparsity occurs when the time- and cost-expensivenessof eliciting human labels reduces the proportion of data that can be labelled, and preference noisearises from uncertainties in the human labelling process (as opposed to perfect oracle labelling).
Wedecided to focus on noise in this work as it is an established way of making oracle experiments morerealistic [28], and also fits in with our discussion of human uncertainties and cognitive biases.
Ourexperiments in Section 4.4 indicate that our methods degrade gracefully in the presence of noise,which gives us some confidence that they will transfer well to human labels.
However, future workshould consider preference sparsity and form, the latter of which will involve modifying the datacollection pipeline and loss function (e.g., to a contrastive loss in the case of pairwise rankings).
Beyond accounting for these three differences whilst still using oracle labels, the next step would beto conduct evaluations using actual human labels.
D Preliminary Experiments on Toy Datasets
In this section, we give detail our preliminary experiments that were run on toy problems to initiallydevelop and validate our approach.
Below we outline the datasets (Section D.1), models and traininghyperparameters (Section D.2), and results (Section D.3) of these experiments.
D.1 Datasets
We introduce three toy datasets, each abstracted from the RL context, to act as benchmark tests forour models.
Each of these datasets uses ordered bags comprised of two-dimensional instances, where here. When (following the method of Section 3.3) an RL agent is trained using a non-Markovian RM model that captures a cognitive bias, should the agent learn to maximise rewards including the bias (which, for example, might lead it to prioritise its peak and final reward rather than seek uniformly good performance), or exluding it (which would revert to uniform prioritisation). This issue of whether intelligent agents should seek to exploit human irrationalities when optimising for their revealed preferences, or appeal to their unbiased “better angels", relates to distinctions between first- and second-order preferences [38] or between experienced and remembered (or decision) utility [7]. We defer this question to those with more relevant expertise but note that regardless of the answer, it is essential to include the biases in the reward model using a method such as the one proposed in each instance has an associated label (called “reward" below, for consistency), and the overall baglabel (return) is the sum of the instance labels.
Toggle Switch
An instance is the position of a toggle switch ts and a value v ; if the switch is on( ts = 1 ), then the reward for the instance is v ; otherwise the reward for the instance is 0 .
Here, thereis no hidden information, as it is possible to calculate the reward for an instance from its contentsts, v alone.
This serves as a null example to elucidate what happens in a Markovian setting.
Push Switch We modify the toggle switch setup so that the instance now represents a push switch( p = 1 if this switch is pressed), where pressing the switch flips a binary hidden state ts (the sameinformation as was previously represented by the toggle switch).
This hidden state then determinesthe reward as before ( v if ts = 0 , else 0 ).
As ts must be tracked between successive instances, itis not possible to determine the reward for an instance solely by observing its contents p, v , so thissetup is non-Markovian.
Dial We generalise the hidden state from a binary switch to a continuous-valued dial.
Given aninstance m, v , the dial’s current value d is moved up or down by m .
The reward is then given as d · v
The problem remains non-Markovian, but now the hidden state that needs tracking, d , is continuousrather than discrete.
D.2 Models and Hyperparameters
When training the MIL models on the toy dataset, we used the Adam optimiser with a batch size ofone (i.e., one bag per batch) to minimise mean squared error (MSE) loss.
Training was performedusing validation loss early stopping, i.e., if the validation loss did not decrease after a certain numberof training epochs (patience value), we terminated the training and selected the model at whichthe validation loss was lowest.
If the patience value was not reached (i.e., the validation loss keptdecreasing), we terminated training after a maximum number of epochs had been reached, and againselected the model at which the validation loss was lowest.
The hyperparameters for training themodels on each dataset (including learning rate (LR) and weight decay (WD)) are given in Table A1.
Dropout was not used.
These hyperparameters were found through a small amount of trial and error,i.e., no formal hyperparameter tuning was carried out.
[Table]
In Tables A2 to A5 we give the architectures for the MIL models we used in the toy datasetexperiments.
The models are a combination of fully connected layers (FC) along with different MILpooling mechanisms.
Some layers also include rectified linear unit (ReLU) activation.
We label thelayers based on the part of the network they belong to: feature extractor (FE), head network (HN), orpooling (P); see Section 3.2.
We also indicate the input and output sizes: b x n indicates an input oroutput where there is a representation of length n for each of the b instances ( b is the size of the inputbag).
[Table] [Table]
[Table] [Table]
D.3 Results
[Table] For each of the toy datasets, we generate 5000 random bags with between 10 and 20 instances perbag (uniformly distributed).
We use an 80/10/10 dataset split for training, validation, and testing,and repeat our experiments with ten different variations of this split (so in total we have ten repeatsof each model type for each dataset).
We show results for both return and reward reconstructionfor the toy datasets in Table A6.
From these results, we can make several observations.
Firstly, asexpected, the ISN architecture only works on the Markovian Toggle Switch dataset, i.e., it fails on thenon-Markovian Push Switch and Dial datasets as it is unable to deal with temporal dependencies.
Wealso note that our two proposed architectures (ISL and CISL) outperform the baseline ESL methodon both return and reward, with the CISL providing the best results overall.
Finally, we observethat a better return performance does not always guarantee better reward performance: for the Dialdataset, the ISL makes better return predictions than the CISL, but worse reward predictions.
Asimilar outcome can be seen for the ESL and the ISN on the Push Switch dataset.
[Table]
E RL Task Details, Data Generation, and Model Hyperparameters
In this section
[Table]
we give more detail on the reward reconstruction experiments for the RL tasks.
First,we give more information about the RL tasks (Section E.1), then explain how we generated datasetsfrom the tasks (Section E.2) [Table] and give the MIL model architectures and hyperparameters used inthe Timer, Moving, Key and Charger tasks (Section E.3)
[Table] and the Lunar Lander task (Section E.4).
Finally, we discuss the limitations of our approach to using MIL RM for the Lunar Lander task(Section E.5).
E.1 Task Details
All five of our experimental RL tasks are depicted in Figure A1. The first four tasks are implemented in Python within a common 2D simulator following the OpenAI Gym standard [9].
The agent’s position x, y is moved by one of five discrete actions: up, down,left, right and no-op.
In the first four cases, the position is moved by 0 .
1 in the specified direction.
The motion vector is then corrupted by zero-mean Gaussian noise with a standard deviation of 0 .
[Figure] in both x and y and clipped into the bounds [0 , 1] 2 .
Zones of interest (spawn zones, treasure, key,charger) are specified as rectangles lying within these bounds.
At time t , the environment state s t(which is directly observed by the MIL RM models) is the 2D vector of the current position [ x t , y t ] ;its dynamics are Markovian given the agent’s chosen action.
The hidden state h t is the task-specificinformation that renders the oracle’s reward function R Markovian:
• Timer: h 0 = 0 and h t +1 =
δ ( h t , s t , a t ) =
h t + 1 ; the hidden state simply tracks the current timestep index.
Reward is given by
[Equation]
s t , a t
[Equation]
h t
[Equation]
[Equation]
= in treasure ( s t ) ·
[Equation]
where in_treasure ([ x t , y t ]) =
if 0 .
4 ≤ x t ≤ 0 .
6 and 0 .
7 ≤ y t ≤ 0 .
9 , and 0 otherwise.
Moving: h 0 =
[0 . 4 , − 0 .
02] , the initial horizontal position (left edge) and velocity of the moving treasure rectangle.
Hidden state dynamics encode the left-right oscillation: h
[Equation]
0 <
[Equation] h
[Equation] h [Equation]
(cid:21) .
Reward is given by R ([ x t , y t ] , a t , h t +1 )
=
if h t +1 ≤ x t ≤ ( h t +1 + 0 . 2) and 0 .
7 ≤ y t ≤
9 , and [-
-]0 otherwise.
• Key: h 0 = 0 , indicating that the agent initialises without the key.
The key collection dynamics are encoded by
[Equation]
Reward is given by R ( s t , a t , h t +1 ) =
in_treasure ( s t ) · h t +1 , where the in_treasure function
is the same as in the Timer task.
• Charger: h 0 = 0 , indicating an initial charge level of zero.
The charging dynamics are encoded by
[Equation]
[Equation]
min( h t + 0 .
if y t ≤ 0 . 3 ,
[Equation]
otherwise.
Reward is given identically to the Key task, R ( s t , a t , h t +1 )
=
in_treasure ( s t ) · h t +1 .
The Lunar Lander task is a modified version of the L UNAR L ANDER C ONTINUOUS - V 2 baseline included as standard in the OpenAI Gym library [9].
We leave the state and action spaces unmodified.
The 8D state vector is [ x, y, v x , v y , θ, ˙ θ, c l , c r ] , where x, y and v x , v y are the landing craft’s horizontal and vertical positions and velocities, θ and ˙ θ are its angle from vertical and angular velocity, and c l ,
c r are two binary contact detectors indicating whether the left and right landing legs are in contact with the ground.
The 2D continuous action [ u m , u s ] is a pair of throttle values for two engines: main u m and side u s .
We also retain the default initialisation conditions (the lander spawns in a narrow zone above the landing pad, with slightly-randomised angle and velocities), the automatic termination of episodes when | x | exceeds 1 (i.e., when the lander leaves the rendered screen area), and the physics that determine how the lander responds to engine activations.
However, we replace the standard reward function with an oracle that rewards the agent for landing on the pad for up to 50 timesteps, and then taking off again to hover within a target zone until an episode time limit ( T = 500 ) is reached.
Rendering this two-stage objective Markovian requires a hidden state h t that tracks the number of timesteps spent on the pad so far.
Formally, reward is given by
[Equation]
if h t +1 < 50 , R no_contact ( s t )
[Equation]
R hover ( s t ) + R shaping ( s t , 1) otherwise, where R pad rewards the agent for being central with both legs on the ground (i.e., on the pad),
R pad ([ x t , y t , v xt , v yt , θ t , ˙ θ t , c lt , c rt ]) = (cid:26) 1 if − 0 .
2 ≤ x t ≤ 0 .
2 and c lt = 1 and c rt = 1 , 0 otherwise,
R no_contact rewards breaking leg-ground contact,
R no_contact ([ x t , y t , v xt , v yt , θ t , ˙ θ t , c lt , c rt ])
= (cid:26) 1 if c lt = 0 and c rt = 0 , 0 otherwise,
R hover rewards aerial positions in a target zone above the pad,
R hover ([ x t , y t , v xt , v yt , θ t , ˙ θ t , c lt , c rt ]) =
(cid:26) 1 if − 0 .
5 ≤ x t ≤ 0 .
5 and 0 .
75 ≤ y t ≤ 1 . 25 , 0 otherwise, and R shaping promotes slow, stable, central flight towards a target vertical position y target ,
R shaping ([ x t , y t , v xt , v yt , θ t , ˙ θ t , c lt , c rt ] , y target ) = 0 .
1 × max (cid:18)
[Equation]
The hidden state dynamics are h t
+1 = (cid:26) min( h t + 1 , 50) if R pad ( s t )
= 1 , h t otherwise, with h 0 = 0 .
E.2 MIL Dataset Generation
We obtain datasets of several thousand trajectories per task, containing a wide distribution of outcomesand return values, as follows.
For each task k , we define a discrete trajectory classifier function
C k : Ξ → C k and a limit p k on the proportion of trajectories in the dataset that are allowed to map toeach class in C k .
These are given as follows:
• Timer: C timer =
num_neg × num_pos, where num_neg = { 0 .. 50 } counts the number of timesteps the agent spends in the treasure while its reward is negative ( t ≤ 50 ), and num_pos = { 0 .. 50 } counts the number while the reward is positive ( t > 50 ).
The number
of classes
is |C timer | = 51 2 = 2601 and the per-class limit is p timer =
• Moving: C moving = num_treasure, where num_treasure = { 0 .. 100 } counts the timesteps spent in the treasure.
|C moving | = 101 and p moving = 0 .
• Key: C key
=
{ no_key , key_no_
• Charger: C charger = num_treasure × charge_bin, where num_treasure = { 0 .. 100 } countsthe timesteps spent in the treasure and charge_bin = { 1 .. 20 } is a binned representation ofthe mean charge level when in the treasure (e.g., 0 . 0 maps to bin 1 , 0 .
48 to bin 10 , 0 .
96 tobin 20 ).
|C charger | = 2020 and p charger = 0 .
• Lunar Lander: C lunar = pad_bin
× take_off × hover_bin, where pad_bin ={ 0 , { 1 .. 49 } , 50+ } is a binned representation of the number of timesteps spent on the landingpad (i.e., zero, fewer than 50 or at least 50 ), take_off = { 0 , 1 } is a binary indicator of whetherthe lander takes off again after being on the pad
and hover_bin = { 0 , { 1 .. 19 } , 20+ } is abinned representation of the number of timesteps spent in the hover zone after being on thepad.
3 |C lunar | = 18 , of which 9 are actually realisable (e
g., the lander cannot take off fromthe pad if it never reached it in the first place) and p lunar = 0 .
For k ∈ { timer , moving , key , charger } , a dataset X k , is assembled iteratively.
On each iteration, wegenerate a length-100 trajectory, ξ , by sampling agent actions uniform-randomly from the actionspace (up, down, left, right, no-op) and running them through the simulator.
Once the trajectory iscomplete, we evaluate its class C k ( ξ ) .
If there are already at least p k × 5000 trajectories in X k withthis class, ξ is discarded.
Otherwise, it is added to X k .
This process repeats until | X k | = 5000 .
The state-action space for Lunar Lander is too large for a random generate-and-select algorithmto terminate in any reasonable time.
Instead, we recycle the length- 500 trajectories generated asa by-product of training the oracle-based RL baselines (black curves in Figure 5, plus six moreruns not included in the figure).
Starting from a bank of 12000 trajectories, filtering based on thethreshold p lunar = 0 .
2 yields a final dataset X lunar with 9762 trajectories.
Although this approachof relying on oracle-trained agents to generate data may initially appear to “put the cart before thehorse", we suggest that it provides a valuable test of the ability of our MIL models to learn fromgoal-directed (c.f. random) trajectories, and thus is a step closer to the online bootstrapping approachof simultaneous RM and RL, which we aim to tackle in future work (see Section 5.3).
E.3 MIL Models and Hyperparameters
The MIL model training on the Timer, Moving, Key, and Charger tasks used the same processas for the toy model training (see Section D.2).
However, we also applied dropout (DO) in thesemodels.
We give the MIL training hyperparameters for each of these tasks in Table A12.
Again, thehyperparameters were found through a small amount of trial and error, i.e., no formal hyperparametertuning was carried out.
[Table]
In Tables A8 to A11 we give the architectures for the MIL models we used in the Timer, Moving,
Key, and Charger tasks.
As in the toy dataset experiments, the models are a combination of fullyconnected layers (FC) along with different MIL pooling mechanisms.
Some layers also include withrectified linear unit (ReLU) activation.
Again, we label the layers based on the part of the networkthey belong to: feature extractor (FE), head network (HN), or pooling (P); see Section 3.2.
We alsoindicate the input and output sizes: b x n an input or output where there is a representationof length n for each of the b instances ( b the size of the input bag)
[Table] [Table]
[Table] [Table]
2) FC + ReLU + DO b x 64 b x3 (FE-3) FC + ReLU + DO b x 32 b x4 (P-1) mil-csc-ins-lstm b x 32 b x5 (HN-1) FC + ReLU + DO b x 34 b x6 (HN-2) FC + ReLU
{++ DO b x 32 b x7 (HN-
 ) FC b x 16 b x8 (P-2) mil-sum b x+}
E.4 Lunar Lander MIL Models and Hyperparameters
There are several differences between the MIL training process for the Lunar Lander task and theother four RL tasks (see Appendix E.3).
Firstly, the reward targets (and as such, return targets) werescaled down by a factor of 100 in order to avoid extremely large gradients from high predictiontargets.
For example, a trajectory with an original return of 700 would have a scaled return of 7.
Secondly, the input data was scaled linearly between -0.5 and 0.5 (using the minimum and maximumrange of each feature).
This was found to give more consistent feature ranges than mean/standarddeviation scaling as was used in the other tasks (this was due to large outliers in certain features,e.g., the rotational features were largely clustered around 0, but had extreme values up to ± 90).
Wegive the training hyperparameters for the lunar lander environment in Table A12.
Again, no formalhyperparameter tuning was carried out, so better performance of these models could potentially beachieved with better parameters (including shorter training times with a higher learning rate).
[Table]
We used similar architectures to the other four RL tasks (see Appendix E.3), but with larger layers;see Tables A13 through A16.
However, the depth of the models remained the same, as did the sizeof the hidden state embedding.
Also note the addition of the Leaky ReLU activation function in thehead networks, which we discuss further in Appendix E.5.
In Table 1, we present results for Lunar Lander using only the top five performing models by reward MSE (50% of all models).
In this section, we discuss why this is the case. When training the Lunar Lander architectures with linear activation functions in the head networks,we found that the models were struggling to learn the correct return predictions as they weremaking negative reward predictions. We know a priori that the Lunar Lander task only has positive
[Table] [Table]
[Table]
[Table]
rewards, therefore we added a Leaky ReLU activation (with a negative slope of 1 × 10 − 6 ) to eacharchitecture’s head network to encourage positive predictions.
This led to an immediate improvementin performance for all model architecture (excluding the Instance Space NN baseline, but this wasexpected to fail on this task as it cannot model temporal dependencies).
However, we found that a proportion of model initialisations remained unable to overcome a certainlocal minimum during training (corresponding to a return prediction MSE of around 2 . 1 ).
In othercases, after this threshold was passed, the model performance would rapidly improve.
Note thisproblem occurred in each of the LSTM-based models (the Instance Space NN architecture was neverobserved to pass this threshold).
Therefore, we focus our evaluation on the top 50% of models, whichis equivalent to discarding the worse performing models, the majority of which had not passed theproblematic threshold during training.
We investigate this training issue further in Figure A2, focusing specifically on the CSC Instance
Space LSTM models.
Although the Leaky ReLU activation encourages models to make positivepredictions, we can see that it does not prevent them entirely.
Furthermore, the models that arenot able to cross the return threshold of 2 . 1 tend to output a greater proportion of negative rewardpredictions.
We thus hypothesise that a better mechanism for preventing negative reward predictionwould increase the chance that a given model training run achieves a return loss of less than 2 . 1 .
Below we list several such mechanisms as alternatives to our current Leaky ReLU approach.
Replace Leaky ReLU with ReLU One option to remove negative reward predictions entirely is touse ReLU rather than Leaky ReLU in the final activation function of the head networks.
However, inthe case that all the reward predictions are negative, the gradient of the network will be zero, so nolearning can take place.
The ability of the network to learn is entirely dependent on achieving at leastone positive prediction in order to generate non-zero gradients, which is determined by the initialnetwork weights.
This could potentially be overcome with different network initialisation approaches,or by simply discarding training runs that fail to begin to learn.
Replace Leaky ReLU with Sigmoid Instead of using a Leaky ReLU or ReLU activation in the headnetwork, the Sigmoid activation function could be used.
This would overcome the gradient issuepresented with the use of ReLU, and would ensure that no negative rewards are predicted.
However, this would require a priori knowledge of the maximum reward target in order to the scale the networkoutputs correctly, and the non-linear activation could make accurate prediction of rewards moredifficult in some cases.
Remove target normalisation A further approach to reduce the number of negative reward predictions would be remove, or at least reduce, the reward target scaling.
As discussed above, thiswas initially included to avoid very large gradients.
However, reducing this scaling would move theaverage reward prediction away from zero, potentially reducing the number of negative predictions.
Care would have to be taken to not reintroduce the large gradient problem.
Regularise reward prediction variance A final alternative approach could be to introduce anadditional loss term that encourages the reward predictions to have low variance.
This would detera large range of reward predictions, leading indirectly to fewer negative reward predictions (see Figure A2).
However, this approach is only applicable to the LSTM models that produce instancepredictions, i.e., it cannot be used with the Embedding Space LSTM network as that architecturedoes not produce reward predictions during training.
Furthermore, this would result in an additionalhyperparameter that requires tuning (a coefficient for the new loss term).
F Further Details on RL Agent Training
Adopting OpenAI
Gym terminology [9], non-Markovian reward functions (both ground-truth oraclesand learnt LSTM-based models) are implemented as wrappers on rewardless base environments.
The role of a wrapper is to track the hidden state of either the oracle or the LSTM throughout anepisode and use this to compute rewards to return to the agent.
In the “Oracle (without hidden state)"baseline, we return the raw environment state (e.g., the 2D position [ x t , y t ] ) to the agent unmodified.
Otherwise, we concatenate the post-update hidden state h t +1 onto the end of the environment state,thereby expanding the state space from the agent’s perspective and making rewards Markovian.
This wrapper-based approach allows us to use a completely vanilla RL algorithm.
For the 2Dnavigation tasks, we use a Deep Q-Network (DQN) agent [31] with the double Q-learning trick [43]enabled.
For Lunar Lander, which has a continuous action space, we use Soft Actor-Critic (SAC)[15].
In both cases we use a value network with ReLU activation functions, which is updated onevery timestep by sampling batches of size 128 from a replay buffer.
Bellman updates use a discount factor of γ = 0 . and are implemented by the Adam optimiser with a learning rate of 1 e − 3 .
A targetnetwork tracks the primary one by Polyak averaging of parameters with a coefficient of 5 e − 3 pertimestep.
Additional hyperparameters are given below:
• 2D tasks (DQN): number of training episodes = 400 ; replay buffer capacity = 5 e 4 ; valuenetwork hidden layer sizes =
[256 , 128 , 64] ; policy greediness ϵ linearly decayed from 1 to0 .
over the first 200 episodes and held constant thereafter. • Lunar Lander (SAC): number of training episodes = 800 ; replay buffer capacity =1 e 5 ; value/policy network hidden layer sizes =
[256 , 256] ; policy entropy regularisation
coefficient α = 0 .
2 ; policy updates with Adam optimiser (learning rate 1 e − 4 ).
G Additional RL Training Analysis for Lunar Lander
[Figure] There is relatively little separation in performance on the shaping reward R shaping and pad contact reward R pad (for the latter, all runs end up reliably achieving the maximum possible reward of 49 , although those using Embedding Space LSTM models take significantly longer to get there). This indicates that all models have been able to recover these components with reasonable fidelity.
However, there are marked differences in performance on R no_contact and R hover (the components relating to the second task stage of taking off and moving to the hover zone).
For R hover , runs using the CSC Instance Space LSTM peak at a return of around 200 from this component, while those using the other two models (almost) never achieve non-zero return, i.e., only the RL agents trained using the CSC Instance Space LSTM RM models reliably learn to hover. This indicates that the model architectures have learnt fundamentally different representations of the reward function and hidden state dynamics, which are effective for policy learning in the case of the CSC model, and highly ineffective for the others.
Also note that, as in the Timer task, runs using CSC Instance Space LSTM models outperform those with direct access to the ground-truth oracle on all components, and especially on R hover . This surprising finding suggests that this model reliably learns hidden state representations that are easier for RL agents to use in policy learning than the ground-truth ones, and potentially that certain errors in the reward prediction may actually be beneficial for the purpose of helping agents to complete the underlying task (in typical RL parlance: the model’s reward function may be better shaped than the ground-truth).
The origins and potential generalisability of this “better-than-oracle" phenomenon warrants significant further investigation.
H Interpretability Plots for Other RL Tasks
In this section, we provided additional interpretability plots for our MIL RM models.
In Section H.1 (Figure A4), we analyse the lunar lander learnt hidden states (like Figure 7). In Section H.2, weprovide the interpretability probes (like Figure 8) for the timer (Figure A5), moving (Figure A6), andkey (Figure A7) tasks.
As before, in these probing plots, we analyse the best performing CISL modelfor each task (according to the reward reconstruction metric).
In both types of plots, we also indicate the hidden state that relates to the agent’s initial position (cyan triangle).
We find that each of the models has managed to learn to separate the three key environment states. They also all capture a binary distinction between the crossover point of t = 50 . However, the CSC Instance Space LSTM model represents this transition in a very sparse manner, i.e, there is a clear gradient between times just below 50 and just above, allowing the RL learning process to more easily identify the crossover point, leading to better RL performance as observed in Figure 5.
As before, we use the best performing repeat for each model as assessed by the reward reconstruction.
E n v i r o n m e n t S t a t e
T r a j e c t
o r y
H i dd e n S t a t e
T r a j e c t
o r y
Optimal (oracle return = 50)
Too Late (oracle return = 41)
Too Early (oracle return = 34)
R e w a r dChallenging (oracle return = 44) t =t =t =t =
t =
[Figure] [Figure] t
[Figure]
50, In t =x t =x x x x x x x
Top : Predicted reward with respect to time and position.
We observe that the model has correctly captured the transition from negative to positive reward in the treasure region at t = 50 , with no reward outside of this region.
Although the reward is positive at t = 50 , the model is uncertain at this point, i.e., the transition from negative to positive reward happens over two timesteps rather than one.
Bottom : Four probes
Optimal : The agent moves into the treasure region at t = 50 and remains there, receiving the maximum possible reward.
Too Late : The agent moves into the treasure region a while after the treasure has already become positive, i.e., it is missing out on reward by not being in the region for as long as possible.
This is reflected in the hidden state plot, where the state transitions from the orange to the red region after the t = 50 boundary.
Too Early : The agent moves into the treasure region before the treasure becomes positive, therefore, while it earns the maximum amount of positive reward, it also earns negative reward, leading to a sub-optimal result.
Challenging :
The agent moves into the treasure region at the correct time, but proceeds to jump in and out of the treasure region before settling, leading to lost reward.
The hidden state trajectories somewhat mimic this movement by transitioning between the orange and red regions, although the jumps are less clear near to the t = 50 transition point, suggesting the model is uncertain at this point.
t =
E n v i r o n m e n t
S t a t e
T r a j e c t
o r y
H i dd e n S t a t e
T r a j e c t
o r y
[Figure]
R e w a r d
Challenging
[Figure]
t
[Figure]
t
[Figure]
t =t =
[Figure]
Top : The predicted reward with respect to time (and as such, position of the treasure region).
The model has learnt to track the treasure region as it moves, although there is noise around the left and right edges of the region, highlighting the difficulty in tracking the treasure region’s horizontal position.
Bottom : Four interpretability probes showing the model’s hidden state transitions.
Note the green dotted region indicates the overall boundary of the treasure region, i.e., the treasure lies somewhere within that boundary, with its true horizontal position dependent on time.
“ x n " labels indicate the agent remaining in a position for n timesteps.
Optimal Left : The agent moves within the treasure region as quickly as possible and then moves with the treasure (keeping within it) for the remainder of the episode.
The hidden state trajectories follow the “In" regions (green and red).
Optimal Right : A similar optimal probe where the agent starts from the right-hand spawn zone rather than the left.
In this case, the agent has further to move before moving into the treasure region, leading to slightly less overall reward than when it starts from the left.
Static : Rather than moving with the treasure region, the agent stays still and allows the treasure to pass over it, gaining reward for some timesteps but not others.
We observe that the hidden state trajectory also reflects this — the state transitions between the “In” and “Out” regions.
Challenging :
The agent takes a while to move towards the treasure region, and then passes in and out of the boundary in which the treasure resides, picking up some reward.
We can see from the hidden state trajectory that this motion is captured in the hidden states — the trajectory transitions between the orange and red regions as the agent passes back and forth through the treasure region.
t =
E n v i r o n m e n t S t a t e
T r a j e c t
o r y
H i dd e n S t a t e
T r a j e c t
o r y
[Figure]
e
[Figure]
a r d
[Figure]
Optimal Left : The agent collects the key as quickly as possible, and then proceeds to move into the treasure region and wait there until the end of the episode.
The hidden state trajectory shows two notable transitions, first when the key is collected orange and then when the enters the
Optimal Right : A similar optimal policy, but from the right-hand spawn zone rather than
the
Note the hidden state trajectory is very similar, apart from the initial hidden state, which is at the opposite side of the blue region, representing the different spawn position.
Suboptimal : The agent passes through the treasure region before collecting the key, and then follows an optimal policy.
We observe a transition in the hidden state trajectory from the blue to the green region that corresponds to the agent’s premature entrance into the treasure region.
Failure case : The agent enters the treasure region without collecting the key and remains there for the rest of the episode.
This highlights a failure case of the model, where the hidden state “drifts” from the green region to the red region, i.e., the model convinces itself that the agent must have picked up the key at some point.
Such causal confusion errors are well-documented in the RM literature [40].
In our case, we suspect that the error is due to a bias in our data generation method inducing a correlation between key possession and time spent in the treasure region.
As described in Appendix E.2, we specifically screen for trajectories where the treasure is visited with the key (the “treasure" class) but not for those where it is visited without it.
There are thus likely to be very few training trajectories that spend a lot of time in the treasure region without collecting the key, making this probe an extreme outlier on which performance is poor.
Adding and selecting for a fourth “key_no_treasure" class during data generation may have mitigated this issue.