Measurement of $\Lambda$ hyperon spin-spin correlations in p+p collisions by
the STAR experiment

Jan Vanek, for the STAR collaboration

Brookhaven National Laboratory

> Polarization of $\Lambda$ hyperons has been observed in various collision
> systems over a wide range of collision energies over the last 50 years since
> its discovery in Fermilab in the 70’s. The existing experimental and
> theoretical techniques were not able to provide a conclusive answer about
> the origin of the polarization. In these proceedings, we discuss the
> possibility to use a new experimental method which utilizes measurement of
> $\Lambda\bar{\Lambda}$, $\Lambda\Lambda$, and $\bar{\Lambda}\bar{\Lambda}$
> pair spin-spin correlations. With this new approach, it should be possible
> to distinguish if the polarization originates from early stage effects, such
> as initial state parton spin correlation, or if it is a final state effect
> originating from hadronization. Furthermore we, study the feasibility to
> perform this measurement in $p+p$ collisions at $\sqrt{s}=200$ GeV collected
> by STAR in 2012 which should provide sufficient statistics of
> $\Lambda\bar{\Lambda}$, $\Lambda\Lambda$, and $\bar{\Lambda}\bar{\Lambda}$
> pairs to perform this measurement.

> PRESENTED AT
>
>
>
>
> DIS2023: XXX International Workshop on Deep-Inelastic Scattering and Related
> Subjects,
> Michigan State University, USA, 27-31 March 2023
>

##  1 Introduction

An interesting discovery by Fermilab was published in 1976. They observed that
$\Lambda$ hyperons produced in $p+$Be collisions with a 300 GeV proton beam
are polarized [1]. This observation is surprising because neither the proton
beam nor the beryllium target was polarized. As a result, experimentalists and
theorists all around the world started investigating this phenomenon.

The $\Lambda$ hyperon polarization can be measured through reconstruction of a
hadronic decay channel $\Lambda^{0}\rightarrow p\pi^{-}$ (and charge
conjugate) and subsequent measurement of the angle ($\theta_{\mathrm{p}}$, or
$\theta^{\star}$) between the decay proton momentum in the $\Lambda$ rest
frame ($p$) and a normal vector to the production plane ($\hat{n}$). This
decay channel is selected, because the proton is preferentially emitted in the
direction of the $\Lambda$ polarization in the $\Lambda$ rest frame. A cartoon
illustrating the production plane determination, using the variables defined
above, is shown in Fig. 1.

Figure 1:  Cartoon illustrating production plane determination. The plane is
defined by momentum of the beam ($\vec{p_{\mathrm{e}}}$) and momentum of
$\Lambda$ ($\vec{p_{\Lambda}}$). Normal vector $\hat{n}$ to the plane is then
given by $\hat{n}=\vec{p_{\mathrm{e}}}\times\vec{p_{\Lambda}}$. The $\Lambda$
polarization is then quantified by measuring the angle $\theta_{\mathrm{p}}$
between the momentum of proton in the rest frame of its mother and $\hat{n}$.
Taken from Ref. [2].

The polarization ($P_{\Lambda}$) is then extracted from the angular
distribution of the protons according to formula

$\frac{\mathrm{d}N}{\mathrm{d}\cos(\theta^{\star})}=1+\alpha
P_{\Lambda}\cos(\theta^{\star}),$ (1)

where $\alpha$ is the weak decay constant of the $\Lambda$ hyperon. This
method was used in the first $\Lambda$ hyperon polarization measurement from
Ref. [1] and other measurements that followed. It is also possible to measure
the $\Lambda$ hyperon polarization with respect to a different reference
direction, e.g. a jet axis, or polarization of the beam for measurements with
polarized beams. A brief overview of experimental results using these
traditional methods is provided in Sec. 2.

Section 3 provides a description of a new method for $\Lambda$ hyperon
polarization measurement, which relies on the determination of $\Lambda$
hyperon pair spin-spin correlation. In addition, the section shows first steps
of analysis utilizing this method in $p+p$ collisions at $\sqrt{s}=200$ GeV by
the STAR experiment.

##  2 Overview of $\Lambda$ hyperon polarization results

One of the results presented in the first $\Lambda$ hyperon polarization paper
is shown in Fig. 2. The polarization $\alpha P_{\Lambda}$ of the $\Lambda$
hyperons produced in $p+$Be collisions with 300 GeV proton beam on a Be
target, rises with the $\Lambda$ transverse momentum ($p_{\mathrm{T}}$).

Figure 2:  The first observation of $\Lambda$ hyperon polarization as a
function of transverse momentum, measured at Fermilab in $p+$Be collisions
with 300 GeV proton beam. Taken from Ref [1].

In pursuit of an explanation of the origin of $\Lambda$ hyperon polarization,
a number of independent investigations were performed over a wide range of
collision systems and energies. A selection of such results is shown in Fig.
3. In this case, the polarization $P_{\Lambda}$ is plotted as a function of
$x_{\mathrm{F}}=p_{\mathrm{z}}^{\Lambda}/p_{\mathrm{beam}}$. The key
observation here is that the polarization appears to depend primarily on the
$x_{\mathrm{F}}$ and not on the collision type or energy.

Figure 3:  Comparison of the $\Lambda$ hyperon polarization as a function of
$x_{\mathrm{F}}$ measured in variety of collision systems and collision
energies. Taken from Ref. [3].

All presented results are from collisions of unpolarized particles. It is also
important to investigate, if the polarization of the produced $\Lambda$
hyperons is correlated with polarization of the beams, for example, in
polarized $\vec{p}+\vec{p}$ collisions at the STAR experiment. An example of
such measurement is presented in Fig. 4 which shows the longitudinal spin
transfer $D_{\mathrm{LL}}$ of $\Lambda$ and $\bar{\Lambda}$ hyperons at
positive pseudorapidity ($0<\eta<1.2$) measured in $\vec{p}+\vec{p}$
collisions at $\sqrt{s}=200$ GeV. No significant longitudinal polarization of
$\Lambda$ hyperons is observed at STAR, which suggests that the beam
polarization does not play a significant role for the polarization of the
$\Lambda$ hyperon at RHIC energies within the studied kinematic range111The
$x_{\mathrm{F}}$ in this $\eta$ region at RHIC is going to be rather small
which will likely lead to small $\Lambda$ polarization, as seen in Fig. 3. It
is not possible to make direct comparison to Fig. 3 due to polarization of the
beams and also a different observable of the measurement..

Figure 4:  Longitudinal spin transfer $D_{\mathrm{LL}}$ of $\Lambda$ (red
circles) and $\bar{\Lambda}$ (blue circles) hyperons at positive
pseudorapidity ($0<\eta<1.2$) measured in $\vec{p}+\vec{p}$ collisions at
$\sqrt{s}=200$ GeV by the STAR experiment. Taken from Ref [5].

The examples above are a small selection of all experimental efforts over the
last 50 years in attempt to explain the origin of $\Lambda$ hyperon
polarization. Unfortunately, none of the measurements, or the theoretical
models, can provide a definitive answer on where the $\Lambda$ hyperon
polarization is generated. In the following section, we investigate a
possibility to improve our knowledge of the phenomenon by measurement of
$\Lambda$ hyperon pair spin-spin correlations.

##  3 $\Lambda$ hyperon spin-spin correlations

In general, most experimental techniques for measurement of $\Lambda$ hyperon
polarization are based on the same general idea. It is the measurement of an
angle, usually denoted $\theta^{\star}$, between a reference direction and
momentum of the decay proton in the hyperon rest frame. The reference
direction can be chosen based on specific physics considerations. As shown
above, the first possibility is to use the production plane. Other common
alternatives are, e.g., the polarization of the beam, or a jet axis, in case
the $\Lambda$ hyperon is part of a jet.

Another possibility is to look for events with two or more $\Lambda$ or
$\bar{\Lambda}$ hyperons and measure the angle $\theta^{\star}_{12}$ between
the decay (anti-)protons, both boosted into the rest frame of their mother
particles. Since these (anti-) protons are preferentially emitted in the
direction of their mother’s polarization, such measurement gives access to
$\Lambda\bar{\Lambda}$, $\Lambda\Lambda$, and $\bar{\Lambda}\bar{\Lambda}$
pair spin-spin correlations [6, 7]. For this method, it is possible to use the
following formula (see also Ref. [7]):

$\frac{\mathrm{d}N}{\mathrm{d}\cos(\theta^{\star}_{12})}=1+\alpha_{1}\alpha_{2}P_{\Lambda_{1}\Lambda_{2}}\cos(\theta^{\star}_{12}),$
(2)

where $\alpha_{1}$ and $\alpha_{2}$ are weak decay constants of $\Lambda$
hyperons in the pair and $P_{\Lambda_{1}\Lambda_{2}}$ is the polarization of
the pair. $\Lambda_{1}$ and $\Lambda_{2}$ can both be either $\Lambda$ or
$\bar{\Lambda}$.

One of the key advantages of this approach is that it should be able to
identify if the polarization comes from initial stage effects, such as spin-
spin correlation of the initial stage partons, or if it is a final state
effect originating from e.g. hadronization. The initial state correlation
should be seen in data as spin-spin correlation of the $\Lambda\bar{\Lambda}$
pairs, as those likely originate from a single $s\bar{s}$ quark pair produced
in the hard partonic scattering. At the same time, no strong correlation is
expected for $\Lambda\Lambda$ and $\bar{\Lambda}\bar{\Lambda}$ pairs, as those
cannot originate from a single $s\bar{s}$ quark pair.

In order to investigate the $\Lambda$ hyperon pair spin-spin correlations in
$p+p$ collisions at $\sqrt{s}=200$ GeV at STAR, it is important to verify that
there are no other known mechanisms to generate non-zero
$P_{\Lambda_{1}\Lambda_{2}}$. This was done with PYTHIA 8.3 simulation of
$p+p$ collisions at $\sqrt{s}=200$ GeV. The extracted
$1/N_{\mathrm{evt}}\mathrm{d}N/\mathrm{d}\cos(\theta^{\star}_{12})$
distributions as a function of $\cos(\theta^{\star}_{12})$ for
$\Lambda\bar{\Lambda}$ and $\Lambda\Lambda$ pairs are shown in Fig. 5. The
distributions are fitted with a linear function which is used to extract the
value of $P_{\Lambda_{1}\Lambda_{2}}$ using equation (2). For both
combinations, the polarization is zero, meaning that pure PYTHIA does not
predict any $\Lambda$ hyperon spin-spin correlations at mid-rapidity in $p+p$
collisions at $\sqrt{s}=200$ GeV.

Figure 5:  $1/N_{\mathrm{evt}}\textrm{d}N/\textrm{d}\cos(\theta^{\star}_{12})$
as a function of $\cos(\theta^{\star}_{12})$ for $\Lambda\bar{\Lambda}$ (left)
and $\Lambda\Lambda$ (right) pairs from PYTHIA 8.3 $p+p$ collisions at
$\sqrt{s}=200$ GeV. The simulation is fitted with a linear function which is
used to extract the polarization $P$ according to equation (2).

The extraction of the $\mathrm{d}N/\mathrm{d}\cos(\theta^{\star}_{12})$
distributions from the data starts with selection of $\Lambda$ and
$\bar{\Lambda}$ hyperon candidates. This is done by pairing protons and pions
reconstructed and identified with the STAR Time Projection Chamber. Each
$\Lambda$ candidate then corresponds to one $p\pi$ pair. Events with two or
more $\Lambda$ candidates were considered for further analysis. For events
which contain at least two such pairs, a 2D distribution is filled where one
axis is the invariant mass of one of the $p\pi$ pairs and the second axis is
the invariant mass of the second $p\pi$ pair.

This is done for two combinations of the $p\pi$ pairs: an unlike-sign (US)
$p\pi$ pair matched with a different US pair from the same event. This
distribution contains both the signal and the combinatorial background. The
background can be estimated using US pairs matched to the like-sign (LS)
$p\pi$ pairs. The US-LS distribution is then subtracted from US-US
distribution and subsequently fitted with a 2D Gaussian function to determine
the $\Lambda$ candidate invariant mass peak mean and width. The signal region
is defined as the mean $\pm 3\sigma$, where both the mean and $\sigma$ are
taken from the fit. Example of two of the 2D invariant mass distributions is
shown in Fig. 6.

Figure 6:  Two-dimensional invariant mass distributions of unlike-sign $p\pi$
pairs matched with unlike-sign $p\pi$ pairs. Left panel shows $p\pi$ pair
charge combinations corresponding to $\Lambda\bar{\Lambda}$ pair candidates,
the right panel shows the $\Lambda\Lambda$ pair candidates.

This procedure is done separately for $\Lambda\bar{\Lambda}$,
$\Lambda\Lambda$, and $\bar{\Lambda}\bar{\Lambda}$ candidate pairs in $p+p$
collisions at $\sqrt{s}=200$ GeV measured by STAR in 2012. The extracted
numbers of signal and background pairs for each of the three possible
combinations is shown in Tab. 1. The number of candidate pairs is going to
provide sufficient statistics to perform this measurement using the 2012 $p+p$
collisions data-set.

$p_{\mathrm{T}}$ [GeV/$c$] | $\Lambda\bar{\Lambda}$ | $\Lambda\Lambda$ | $\bar{\Lambda}\bar{\Lambda}$
---|---|---|---
$N_{\mathrm{sig}}$ | $N_{\mathrm{bkg}}$ | $N_{\mathrm{sig}}$ | $N_{\mathrm{bkg}}$ | $N_{\mathrm{sig}}$ | $N_{\mathrm{bkg}}$
$\Lambda_{1}$ $0.5-5.0$ | 26767 | 6929 | 7090 | 3593 | 4670 | 2495
$\Lambda_{2}$ $0.5-5.0$
Table 1: Number of $\Lambda\bar{\Lambda}$, $\Lambda\Lambda$, and
$\bar{\Lambda}\bar{\Lambda}$ signal candidate and background pairs in $p+p$
collisions at $\sqrt{s}=200$ GeV measured by STAR in 2012.

##  4 Summary

The $\Lambda$ hyperon polarization puzzle is one of the main unresolved
mysteries of the experimental high energy particle physics. The polarization
has been observed in several different collision systems at various energies.
The magnitude of the polarization appears to depend primarily on
$x_{\mathrm{F}}$ and not much on the specific collision energy. Despite
enormous experimental and theoretical efforts to explain the $\Lambda$ hyperon
polarization, no conclusive answer was found. In order to improve the
knowledge in this field, a new method was developed which relies on
measurement of $\Lambda$ hyperon pair spin-spin correlations. Any non-zero
signal measured in $p+p$ collisions at $\sqrt{s}=200$ GeV at STAR would
provide more insight into $\Lambda$ hyperon polarization, as a simulation
using PYTHIA 8.3 predicts no spin-spin correlation signal. The number of
extracted $\Lambda\bar{\Lambda}$, $\Lambda\Lambda$, and
$\bar{\Lambda}\bar{\Lambda}$ candidate pairs extracted from the aforementioned
STAR data-set is sufficient to perform this type of measurement and thus is
going to provide important additional insight into $\Lambda$ hyperon
polarization physics in $p+p$ collisions at RHIC energies.

## References

  * [1] Bunce G., _et al._ , $\Lambda^{0}$ Hyperon Polarization in Inclusive Production by 300-GeV Protons on Beryllium, Phys. Rev. Lett. 36, 1113 (1976)
  * [2] Airapetian A., _et al._ [HERMES Collaboration], Transverse Polarization of $\Lambda$ and $\bar{\Lambda}$ Hyperons in Quasireal Photoproduction, Phys.Rev.D 76, 092008 (2007)
  * [3] Aad G., _et al._ [ATLAS Collaboration], Measurement of the transverse polarization of $\Lambda$ and $\bar{\Lambda}$ hyperons produced in proton-proton collisions at $\sqrt{s}=7$ TeV using the ATLAS detector, Phys. Rev. D 91, 032004 (2015)
  * [4] Guan Y., _et al._ [BELLE Collaboration], Observation of Transverse $\Lambda/\bar{\Lambda}$ Hyperon Polarization in $e^{+}e^{-}$ Annihilation at Belle, Phys. Rev. Lett. 122, 042001 (2019)
  * [5] Adam J., _et al._ [STAR Collaboration], Improved measurement of the longitudinal spin transfer to $\Lambda$ and $\bar{\Lambda}$ hyperons in polarized proton-proton collisions at $\sqrt{s_{\mathrm{NN}}}=200$ GeV, Phys. Rev. D 98, 112009 (2018)
  * [6] Törnqvist, N.A. Suggestion for Einstein-Podolsky-Rosen experiments using reactions like $e^{+}e^{-}\rightarrow\Lambda\bar{\Lambda}\rightarrow\pi^{-}p\pi^{+}\bar{p}.$, Found Phys 11, 171–177 (1981)
  * [7] Gong, W., _et al._ , Measurement of Bell-type inequalities and quantum entanglement from $\Lambda$-hyperon spin correlations at high energy colliders, Phys. Rev. D 106, L031501 (2022)


# Bayesian Inverse Reinforcement Learning for Non-Markovian Rewards

Noah Topper, Alvaro Velasquez, George Atia11footnotemark: 1 University of
Central FloridaUniversity of Colorado Boulder

###### Abstract

Inverse reinforcement learning (IRL) is the problem of inferring a reward
function from expert behavior. There are several approaches to IRL, but most
are designed to learn a Markovian reward. However, a reward function might be
non-Markovian, depending on more than just the current state, such as a reward
machine (RM). Although there has been recent work on inferring RMs, it assumes
access to the reward signal, absent in IRL. We propose a Bayesian IRL (BIRL)
framework for inferring RMs directly from expert behavior, requiring
significant changes to the standard framework. We define a new reward space,
adapt the expert demonstration to include history, show how to compute the
reward posterior, and propose a novel modification to simulated annealing to
maximize this posterior. We demonstrate that our method performs well when
optimizing according to its inferred reward and compares favorably to an
existing method that learns exclusively binary non-Markovian rewards.

## Introduction

Inverse reinforcement learning (IRL) has become an area of increasing interest
in the last decade. Given an expert demonstration, the goal is to infer the
reward function that the expert is optimizing. We are typically given a
reward-free Markov Decision Process (MDP) and a demonstration in the form of a
set of state-action pair trajectories. We then infer a reward function over
the MDP.

IRL is fundamentally an ill-posed problem; there are many possible reward
functions that yield the same optimal behavior. As such, there are various
methods for settling this ambiguity, including margin optimization (Ng and
Russell 2000) (Abbeel and Ng 2004) (Neu and Szepesvari 2012), maximum entropy
(Ziebart et al. 2008) (Wulfmeier, Ondruska, and Posner 2015) (Boularias,
Krömer, and Peters 2012), and Bayesian (Ramachandran and Amir 2007) (Lopes,
Melo, and Montesano 2009) (Levine, Popovic, and Koltun 2011) methods. We focus
on the Bayesian approach.

There has also been recent interest in non-Markovian RL using reward machines
(RMs) (Toro Icarte et al. 2020). In such settings, the RL agent must optimize
for a non-Markovian reward, meaning one that is history-dependent. RMs are a
certain finite-state machine representation of such memory-based reward
functions.

There has even been recent work in _inferring_ these RMs by observing the
environment (Icarte et al. 2021) (Xu et al. 2019) (Xu et al. 2020). However,
such methods must be able to observe the reward signal. We propose a method
for inferring an RM from expert behavior _only_.

In this paper, we extend the standard Bayesian IRL (BIRL) framework to the
non-Markovian setting using RMs and develop a novel simulated annealing
solution. We then evaluate our algorithm on a testbed of gridworlds and
compare against an exiting algorithm which learns exclusively binary non-
Markovian rewards.

## Preliminaries

###### Definition 1

A _reward-free Markov Decision Process $M$_ is a tuple $(S,s_{I},A,T)$ with
$S$ a finite set of states, $s_{I}$ the initial state, $A$ a finite set of
actions, and $T:S\times A\times S\to[0,1]$ a probabilistic transition
function.

In standard IRL, we are given MDP $M$ and demonstration $\mathcal{D}$, and
must infer a reward function $R:S\to\mathbf{R}$. We assume $\mathcal{D}$ is
produced by some expert, and we wish to infer an $R$ that best explains its
behavior.

In the BIRL framework, we define a prior $P(R)$ over possible reward functions
and $\mathcal{D}$ is given as a set of $N$ state-action pairs,
$\mathcal{D}=\\{(s_{i},a_{i})\mid 1\leq i\leq N\\}$. Let $Q(s,a,R)$ denote the
optimal Q-value of $(s,a)$ in $M$ given $R$. We then assume the expert behaves
according to the Boltzmann rationality model: $P(a\mid s,R)=e^{\alpha
Q(s,a,R)}/\Sigma_{a^{\prime}}e^{\alpha Q(s,a^{\prime},R)}$.

This is a sort of softmax over action values with rationality factor $\alpha$,
where actions with higher long-term reward are exponentially more likely. The
higher $\alpha$ is, the closer the expert is to optimal. Assuming
$\mathcal{D}$ is a single trajectory, we can compute its probability as:

$\begin{split}P(\mathcal{D}\mid R)&=\prod_{i}P(a_{i}\mid s_{i},R)P(s_{i}\mid
s_{i-1},a_{i-1})\\\ &=\prod_{i}\frac{e^{\alpha
Q(s_{i},a_{i},R)}}{\Sigma_{a}e^{\alpha Q(s_{i},a,R)}}\prod_{i}P(s_{i}\mid
s_{i-1},a_{i-1})\\\
&=\frac{e^{\alpha\Sigma_{i}Q(s_{i},a_{i},R)}}{\Pi_{i}\Sigma_{a}e^{\alpha
Q(s_{i},a,R)}}\prod_{i}P(s_{i}\mid s_{i-1},a_{i-1}).\\\ \end{split}$ (1)

The product term on the right is simply the probability of observing the
demonstrated state sequence, given the action sequence. This is determined by
the dynamics and does not vary with $R$, so we can treat it as a constant. If
$\mathcal{D}$ contains multiple trajectories, only this right term will be
affected, and will still be constant with respect to $R$. Thus we can in
general say: $P(\mathcal{D}\mid R)\propto
e^{\alpha\Sigma_{i}Q(s_{i},a_{i},R)}/\Pi_{i}\Sigma_{a}e^{\alpha
Q(s_{i},a,R)}$.

We can then compute the reward posterior as $P(R\mid\mathcal{D})\propto
P(\mathcal{D}\mid R)P(R)$, and we can find the likelihood ratio between any
two reward functions, which is all we need.

Ramachandran and Amir (2007) showed that the mean of the reward posterior is a
good estimator, in the sense that optimizing with respect to it will minimize
expected policy loss. They use Markov chain Monte Carlo (MCMC) to sample from
the posterior and average many such samples.

To estimate the shape of a distribution using MCMC, start at a random point,
generate a random neighbor, transition to this neighbor according to some
probabilistic rule, and repeat. With a well-chosen rule, we will visit each
point proportional to its posterior probability.

Ramachandran and Amir (2007) discretize the space of reward vectors and
transition from $R_{1}$ to $R_{2}$ with probability:
$\min\left(1,\frac{P(R_{2}\mid\mathcal{D})}{P(R_{1}\mid\mathcal{D})}\right)=\min\left(1,\frac{P(\mathcal{D}\mid
R_{2})P(R_{2})}{P(\mathcal{D}\mid R_{1})P(R_{1})}\right)$. If $R_{2}$ is more
likely than $R_{1}$, we transition with probability $1$. If not, we still want
to accept it sometimes to explore, so we move with probability equal to their
likelihood ratio.

## Methodology

### Non-Markovian BIRL Framework

In our setting, the reward is a function of histories rather than MDP states.
We label our MDP states with relevant observations, and the reward is a
function over label sequences.

###### Definition 2

A _labeled reward-free Markov Decision Process $M$_ is a tuple
$(S,s_{I},A,T,\Sigma,L)$ with $S$ a finite set of states, $s_{I}$ the initial
state, $A$ a finite set of actions, $T:S\times A\times S\to[0,1]$ a
probabilistic transition function, $\Sigma$ a finite alphabet, and
$L:S\to\Sigma$ a state-labeling function.

###### Definition 3

A _reward machine_ $R$ is a tuple $(Y,y_{I},\Sigma,\Gamma,\tau,\varrho)$ with
$Y$ a finite set of states, $y_{I}$ the initial state, $\Sigma$ a finite input
alphabet, $\Gamma\subset\mathbf{R}$ a finite set of rewards,
$\tau:Y\times\Sigma\to Y$ a transition function, and
$\varrho:Y\times\Sigma\to\Gamma$ a reward function.

Note that we are assuming a finite set of possible reward values. We will take
this set as known during inference.

###### Definition 4

Given a labeled reward-free MDP $M=(S,s_{I},A,T,\Sigma,L)$ and RM
$R=(Y,y_{I},\Sigma,\Gamma,\tau,\varrho)$, the _product_ $M\times R$ is a
standard MDP $(S^{\prime},s_{I}^{\prime},A,T^{\prime},R^{\prime})$ where:

  * •

$S^{\prime}=S\times Y$

  * •

$s_{I}^{\prime}=(s_{I},y_{I})$

  * •

$T^{\prime}((s,y),a,(s^{\prime},y^{\prime}))=\\\
\begin{cases}T(s,a,s^{\prime})&\text{if }\tau(y,L(s^{\prime}))=y^{\prime},\\\
0&\text{otherwise}\end{cases}$

  * •

$R^{\prime}((s,y),a,(s^{\prime},y^{\prime}))=\varrho(y,L(s^{\prime}))$

Intuitively, we take $M$ and $R$ together to get a standard MDP. For $M\times
R$, we take actions in $M$ and feed the resulting labels into $R$ to generate
rewards. Rewards then depend only on the joint state of the system and are
thus Markovian. We assume the expert is optimizing such a joint MDP, but the
learner only has access to $M$ and $\mathcal{D}$. To estimate $R$, we must
depart from the classic BIRL framework in a number of important ways.

Firstly, $\mathcal{D}$ must include more than just state-action pairs. We are
inferring a non-Markovian reward function, so knowing which actions the expert
took from which MDP states is insufficient; we also need to know the relevant
history.

Let $\mathcal{D}=\\{(s_{i},\lambda_{i},a_{i})\mid 1\leq i\leq N\\}$, where
$\lambda_{i}\in\Sigma^{*}$ is the label sequence the expert observed up to the
point of taking action $a_{i}$. Given RM $R$, we can take the product $M\times
R$ to get a standard MDP. We then project each $(s_{i},\lambda_{i},a_{i})$
onto the state space of $M\times R$, where $\lambda_{i}$ induces a specific
state in $R$ by stepping through the hypothesis according to $\tau$.

Thus for each hypothesis $R$ we have a standard MDP and a set of state-action
pairs $\mathcal{D}$. We can now use BIRL to compute the posterior for $R$. In
$M\times R$, let $Q(s,\lambda,a,R)$ denote the optimal Q-value of taking
action $a$ from state $s$, having observed trace $\lambda$. The posterior is
computed similar to before, with $P(R\mid\mathcal{D})\propto P(\mathcal{D}\mid
R)P(R)$, where now $P(\mathcal{D}\mid R)\propto
e^{\alpha\Sigma_{i}Q(s_{i},\lambda_{i},a_{i},R)}/\Pi_{i}\Sigma_{a}e^{\alpha
Q(s_{i},\lambda_{i},a,R)}$.

Next, even if we can sample RMs from this posterior, it is not clear what it
would mean to average them. Instead, we seek a MAP estimator. Simulated
annealing is a method similar to MCMC, but which seeks to maximize some
distribution, rather than estimate its whole shape. As such, it naturally fits
our setting. Let us summarize the process.

Let $\mathcal{R}_{n}(\Sigma,\Gamma)$ be the space of RMs with $n$ states over
alphabet $\Sigma$ and reward set $\Gamma$. Let $m=|\Sigma|$. An element of
$\mathcal{R}_{n}$ is a tuple $(Y_{n},t,r)$ where $Y_{n}=\\{1,...,n\\}$ is the
state space (assume $1$ is the initial state), $t\in Y^{n\times m}_{n}$ is a
transition matrix, and $r\in\Gamma^{n\times m}$ is a reward matrix.

Fix $n$ as known. Choose an initial point $R_{1}$ by setting the matrices
uniformly at random. Then iteratively propose a “neighboring” RM $R_{2}$ and
transition to it with probability: $\min\left(1,\left(\frac{P(\mathcal{D}\mid
R_{2})P(R_{2})}{P(\mathcal{D}\mid R_{1})P(R_{1})}\right)^{1/T_{i}}\right)$,
with “temperature” $T_{i}$.

This is equivalent to simulated annealing with energy equal to the log
posterior of the RMs. With the proper temperature schedule and choice of
neighbors, we will in principle eventually converge to the peak of the
distribution.

Now we consider how to propose neighbors. In early experiments, we tried
perturbing a single transition, altering one entry of either the $t$ or $r$
matrix to a new random state or reward. This proved to be too little change.
We often need to change multiple factors at once to escape a local optimum.

Instead, fix a probability $p$. Each hypothesis transition (i.e. each entry of
$t$ and $r$) is, with independent probability $p$, changed to a new uniformly
random value, with a minimum of one change. In fact, we use a probability
$p_{i}$, gradually decaying the value along with the temperature. This allows
us to make large changes early on but gradually slow down as we narrow in on a
better region of the hypothesis space.

As a technical detail, we restrict attention to RMs in which all states are
reachable from the start state. This prevents “repeat” hypotheses from showing
up at different sizes. We also consider the “trivial” RM, one with all zero
rewards, to be invalid. If such an RM is sampled, we simply resample.

Before moving on, there is an issue in choosing $\alpha$. If it is set too
low, certain hypotheses that are actually less consistent with the data will
look more plausible than the true RM. We must set $\alpha$ sufficiently high
to assign most of our probability mass on a good region of reward space.
However, if $\alpha$ is set too high, this makes it difficult to move through
reward space. Even rather similar hypotheses may have extremely different
likelihoods, trapping us at local optima.

This is where using temperature helps. By raising the likelihood ratio to
$1/T_{i}$, we “smooth out” the differences between neighboring points.
However, the likelihood ratio $P(\mathcal{D}\mid R_{2})/P(\mathcal{D}\mid
R_{1})$ still takes on extreme values for large $\alpha$, making the prior
ratio $P(R_{2})/P(R_{1})$ tiny by comparison, nullifying the utility of any
prior information. We thus propose a modified transition rule, accepting
$R_{2}$ over $R_{1}$ with probability:
$\min\left(1,\left(\frac{P(\mathcal{D}\mid R_{2})}{P(\mathcal{D}\mid
R_{1})}\right)^{1/T_{i}}\frac{P(R_{2})}{P(R_{1})}\right)$. We can now set
$\alpha$ high, use $T_{i}$ to get a more relaxed distribution, and maintain
the importance of the prior information.

Repeat this process for $N$ iterations and output the best hypothesis seen. We
choose the $R$ with the largest “score” $P(R\mid\mathcal{D})^{1/T}P(R)$,
smoothed by temperature $T$. We must use the same $T$ for all $R$ to make a
consistent comparison, rather than using $T_{i}$ at the point where $R$ is
sampled. We use the minimum value of the temperature schedule for $T$.

### Setting the Parameters and Schedules

Setting $\alpha$ requires some tuning. If during trial runs the chosen RM $R$
does not fit the data very well, $\alpha$ should be raised. However, if
$\alpha$ is set too high, there could be difficulty moving through reward
space. The higher $\alpha$ is, and the more data we generate for the
demonstration, the more extreme $P(\mathcal{D}\mid R_{2})/P(\mathcal{D}\mid
R_{1})$ will tend to be. Accordingly, a higher initial temperature $T_{0}$
will need to be chosen, with a more gradual declining temperature schedule, to
move smoothly through reward space.

We set our temperature schedules as follows. Set $T_{0}$ to some large initial
value. Every $k$ steps, decay the temperature by constant factor
$0<\beta_{T}<1$. Repeat until we hit minimum temperature $T$, at which point
the temperature stays constant. $T$ should be chosen in conjunction with
$\alpha$ such that we get a “reasonable“ distribution over hypothesis space.
Meaning, similar hypotheses usually do not have posteriors that differ by
orders of magnitude, and prior information is still sometimes useful in
deciding between hypotheses.

We follow a similar pattern for the perturbance probability schedule. Set
$p_{0}$ to some middling initial value. Every $k$ steps, decay by constant
factor $0<\beta_{p}<1$ to a minimum value of $p$, chosen so that some small
number of transitions are still altered at each step (around $1$ or $2$ on
average).

### Setting the Prior

Finally, we must set the prior $P(R)$. In early experiments, we simply used a
uniform prior. However, we found it useful to place a prior in favor of
“simple” transitions. That is, a prior in favor of zero-reward, self-
transitions.

For each transition, suppose it has prior probability $p_{r}$ of being zero-
reward and $p_{t}$ of being a self-transition. The rest of the probability
mass is split uniformly between the remaining possible rewards and states.
Suppose $G=|\Gamma|$. If $R$ has $i$ zero-reward transitions and $j$ self-
transitions, then its prior is
$\tilde{P}(R)=p_{r}^{i}[(1-p_{r})/(G-1)]^{nm-i}p_{t}^{j}[(1-p_{t})/(n-1)]^{nm-j}$.
Note that this is an _unnormalized prior_ ; it does not take into account how
many different RM structures there are. However, it suffices to calculate the
prior ratio, which is all that is needed. Thus,
$P(R_{2})/P(R_{1})=\tilde{P}(R_{2})/\tilde{P}(R_{1})$.

## Related Work

As mentioned in the introduction, approaches to IRL include margin
optimization (Ng and Russell 2000) (Abbeel and Ng 2004) (Neu and Szepesvari
2012), maximum entropy (Ziebart et al. 2008) (Wulfmeier, Ondruska, and Posner
2015) (Boularias, Krömer, and Peters 2012), and Bayesian (Ramachandran and
Amir 2007) (Lopes, Melo, and Montesano 2009) (Levine, Popovic, and Koltun
2011) methods. All of these infer a Markovian reward function. Recent work
considered non-Markovian RL using RMs (Toro Icarte et al. 2020), including
inferring such RMs (Icarte et al. 2021) (Xu et al. 2019) (Xu et al. 2020).
However, these methods require access to the reward signal, which we do not
assume.

More closely related are the works of (Memarian et al. 2020) and (Vazquez-
Chanlatte et al. 2017). In the former, the authors develop a deep maximum
entropy IRL algorithm for “temporally extended tasks”, specified by a
deterministic finite automaton (DFA). A series of observed “subgoals” is
either accepted as satisfying the overall goal, or otherwise rejected. They
apply the $L^{*}$ algorithm (Angluin 1987) to learn the DFA structure and then
infer a reward function over the joint states of the product MDP/DFA. However,
this work operates with a functionally binary reward ($1$ for achieving the
goal, $0$ otherwise). Then, a “teacher” answers the learner’s queries about
which sequences satisfy the goal. This implicitly gives reward information to
the agent. In contrast, our algorithm applies with no reward information. Our
algorithm can also learn non-binary rewards.

In the latter work, the authors also learn a task specification, which they
interpret as a Boolean non-Markovian reward. In contrast to the former work,
they manage to do this purely from expert behavior. Their method has a lot of
complexity to it, and in any case still only learns binary rewards, in
contrast to our method.

## Experiments

We now test our algorithm on a set of gridworlds. Memarian et al. (2020)
showed that applying standard memoryless IRL when the true reward signal is
non-Markovian gives poor results, often resulting in no reward at all.
Therefore, a more sophisticated method is certainly necessary.

As mentioned before, the method in (Memarian et al. 2020) does not apply in
our setting, since it requires more information than just expert behavior.
Vazquez-Chanlatte et al. (2017) is the only work we know of which solves
roughly the same problem we do, to infer a non-Markovian reward directly from
the expert demonstration. As such, it may seem natural to compare against
their method. In fact, we use the running example from their paper as one of
our gridworlds.

However, this latter method is prohibitively slow in our setting. It requires
first defining a “concept class”, i.e. the hypothesis space. Just to set up
the problem, it is necessary to determine the subset relationships between
_all pairs of hypotheses_ , in terms of the language accepted by each binary
specification (of course, many pairs are incomparable). A lattice structure is
then built from the hypothesis space and operated on to find the best
specification. Vazquez-Chanlatte et al. (2017) do this in a reasonable amount
of time by considering a concept space of only 930 specifications. In our
setting, we allow all possible RMs of a fixed size and end up with millions of
specifications, even for the small problems we are considering. This results
in trillions of pairs, which takes weeks or months to deal with. As such, this
method is not suitable for our setting.

We consider three gridworlds. The first is from (Vazquez-Chanlatte et al.
2017). There, the expert attempts to reach a recharging station while avoiding
lava and water. If the expert does get wet, it can navigate to a towel to dry
off before recharging. The second is from (Toro Icarte et al. 2020), in which
the expert attempts to deliver a cup of coffee to the office, while avoiding
any decorations. Finally, we create a modification of the coffee gridworld
where one of the machines produces weak coffee, yielding a reward of $1$ when
brought to the office, and the other produces strong coffee, yielding a reward
of $2$.

In all cases, we set $\alpha=50$. We also set the transition prior as
$p_{t}=3/5$ and the reward prior as $p_{r}=3/4$. The values of the other
parameters are shown in Table 1. In each example, we perform three independent
runs of the algorithm and keep the highest scoring hypothesis. We then train
the learning agent on this inferred reward function using policy iteration and
finally run the learned policy on the given gridworld for $100$ episodes. In
each episode, we sum the rewards achieved by the agent at each step, then
compute an average over the $100$ episodes. We do the same for the expert and
compare. The final results can be seen in Table 1.

For both clarity and reminder, in Table 1, $n$ denotes how many states are in
the true reward machine (and is thereby the size of each hypothesis
considered), “runs” and “ep_len” are the number and length of episodes used to
generate the demonstration, respectively, $N$ is the number of samples we
generate in the simulated annealing process, and $r_{e}$ and $r_{a}$ are the
average expert and agent reward, respectively.

The learner performs as well as or better than the expert on these gridworlds.
The learner can do better because we make the expert noisy, while we allow the
agent to fully optimize the inferred reward function. If the agent is able to
infer an accurate reward function from the noisy expert, it can outperform the
given demonstration.

| Recharge | Coffee | Multi Coffee
---|---|---|---
$n$ | 3 | 3 | 4
runs | 1,000 | 100 | 300
ep_len | 25 | 100 | 100
$N$ | 2,000 | 1,000 | 10,000
$T_{0}$ | 500,000 | 100,000 | 1,000,000
$T$ | 200 | 300 | 50
$\beta_{T}$ | 0.98 | 0.96 | 0.99
$p_{0}$ | 0.5 | 0.5 | 0.5
$p$ | 1/16 | 1/12 | 1/16
$\beta_{p}$ | 0.99 | 0.99 | 0.995
$k$ | 5 | 5 | 10
$r_{e}$ | 1.0 | 0.55 | 1.37
$r_{a}$ | 1.0 | 0.68 | 2.0
Table 1: Experimental Results

## Conclusion and Future Work

We have devised an algorithm for extending existing BIRL techniques to the
domain of non-Markovian reward functions. We define our hypothesis space to be
the space of RMs, project our demonstration onto the state space of each
hypothesis product MDP $M\times R$, and apply simulated annealing to find a
MAP estimator for the posterior distribution. There are several avenues for
future work.

First, concerning the algorithm, there could be a more efficient optimization
method than simulated annealing. Other variants could make use of simulated
tempering (Marinari and Parisi 1992), parallel tempering (Earl and Deem 2005),
beam search, or any other discrete optimization method.

Further, since we only need Q-values for the state-action pairs found in
$\mathcal{D}$, it would be more efficient to obtain an estimate of the Q
function only on $\mathcal{D}$. Potential algorithms include prioritized
sweeping, trajectory sampling, and real-time dynamic programming. It is not
immediately obvious, however, how they would map onto the problem at hand.

Finally, there are methods introduced in (Michini and How 2012) for improving
the efficiency of BIRL. A kernel function is introduced to measure similarity
between states. When updating the guess for the reward vector, states are
sampled according to their similarity to those in the demonstration. While
these methods could improve efficiency, we need to understand how to apply
such methods with the expansion of our state space.

For conceptual advancements, our method is limited by assuming the number of
states has a known fixed value of $n$. It would be preferable if $n$ was
itself a learned parameter. However, this would make the sizes of our reward
and transition matrices change as $n$ changes. Thus, the dimensionality of our
hypothesis space would change as we move through it. The method of “reversible
jump” MCMC (Green 1995) allows both the number of parameters and their values
to change, which could possibly be adapted to our current setting.

Finally, the entire method is a passive one. There has been work in using
active methods, where the learner issues queries to the expert, to enhance
learning (Lopes, Melo, and Montesano 2009). The idea is to query states for
which the learner has maximum entropy, and thus will learn the most by seeing
the expert’s action. An avenue of future investigation would be to extend
these methods to the current setting.

## Acknowledgments

This work was supported in part by NSF Award CCF-2106339 and NSF CAREER Award
CCF-1552497.

## References

  * Abbeel and Ng (2004) Abbeel, P.; and Ng, A. Y. 2004.  Apprenticeship Learning via Inverse Reinforcement Learning.  In _Proceedings of the Twenty-first International Conference on Machine Learning_. ACM Press.
  * Angluin (1987) Angluin, D. 1987.  Learning regular sets from queries and counterexamples.  _Information and computation_ , 75(2): 87–106.
  * Boularias, Krömer, and Peters (2012) Boularias, A.; Krömer, O.; and Peters, J. 2012.  Structured Apprenticeship Learning.  In Flach, P. A.; De Bie, T.; and Cristianini, N., eds., _Machine Learning and Knowledge Discovery in Databases_ , 227–242. Berlin, Heidelberg: Springer Berlin Heidelberg.
  * Earl and Deem (2005) Earl, D. J.; and Deem, M. W. 2005.  Parallel tempering: Theory, applications, and new perspectives.  _Physical Chemistry Chemical Physics_ , 7(23): 3910.
  * Green (1995) Green, P. 1995.  Reversible jump Markov chain Monte Carlo computation and Bayesian model determination.  _Biometrika_ , 82(4): 711–732.
  * Icarte et al. (2021) Icarte, R. T.; Waldie, E.; Klassen, T. Q.; Valenzano, R.; Castro, M. P.; and McIlraith, S. A. 2021.  Learning Reward Machines: A Study in Partially Observable Reinforcement Learning.
  * Levine, Popovic, and Koltun (2011) Levine, S.; Popovic, Z.; and Koltun, V. 2011.  Nonlinear Inverse Reinforcement Learning with Gaussian Processes.  In Shawe-Taylor, J.; Zemel, R.; Bartlett, P.; Pereira, F.; and Weinberger, K., eds., _Advances in Neural Information Processing Systems_ , volume 24. Curran Associates, Inc.
  * Lopes, Melo, and Montesano (2009) Lopes, M.; Melo, F.; and Montesano, L. 2009.  Active Learning for Reward Estimation in Inverse Reinforcement Learning.
  * Marinari and Parisi (1992) Marinari, E.; and Parisi, G. 1992.  Simulated Tempering: A New Monte Carlo Scheme.  _Europhysics Letters (EPL)_ , 19(6): 451–458.
  * Memarian et al. (2020) Memarian, F.; Xu, Z.; Wu, B.; Wen, M.; and Topcu, U. 2020.  Active Task-Inference-Guided Deep Inverse Reinforcement Learning.  In _2020 59th IEEE Conference on Decision and Control (CDC)_. IEEE.
  * Michini and How (2012) Michini, B.; and How, J. P. 2012.  Improving the efficiency of Bayesian inverse reinforcement learning.  In _2012 IEEE International Conference on Robotics and Automation_ , 3651–3656.
  * Neu and Szepesvari (2012) Neu, G.; and Szepesvari, C. 2012.  Apprenticeship Learning using Inverse Reinforcement Learning and Gradient Methods.
  * Ng and Russell (2000) Ng, A. Y.; and Russell, S. 2000.  Algorithms for Inverse Reinforcement Learning.  In _in Proc. 17th International Conf. on Machine Learning_ , 663–670. Morgan Kaufmann.
  * Ramachandran and Amir (2007) Ramachandran, D.; and Amir, E. 2007.  Bayesian Inverse Reinforcement Learning.  In _Proceedings of the 20th International Joint Conference on Artifical Intelligence_ , IJCAI’07, 2586–2591. San Francisco, CA, USA: Morgan Kaufmann Publishers Inc.
  * Toro Icarte et al. (2020) Toro Icarte, R.; Klassen, T. Q.; Valenzano, R.; and McIlraith, S. A. 2020.  Reward Machines: Exploiting Reward Function Structure in Reinforcement Learning.  _J. Artif. Int. Res._ , 73.
  * Vazquez-Chanlatte et al. (2017) Vazquez-Chanlatte, M.; Jha, S.; Tiwari, A.; Ho, M. K.; and Seshia, S. A. 2017.  Learning Task Specifications from Demonstrations.  _arXiv_.
  * Wulfmeier, Ondruska, and Posner (2015) Wulfmeier, M.; Ondruska, P.; and Posner, I. 2015.  Maximum Entropy Deep Inverse Reinforcement Learning.
  * Xu et al. (2019) Xu, Z.; Gavran, I.; Ahmad, Y.; Majumdar, R.; Neider, D.; Topcu, U.; and Wu, B. 2019\.  Joint Inference of Reward Machines and Policies for Reinforcement Learning.
  * Xu et al. (2020) Xu, Z.; Wu, B.; Ojha, A.; Neider, D.; and Topcu, U. 2020.  Active Finite Reward Automaton Inference and Reinforcement Learning Using Queries and Counterexamples.
  * Ziebart et al. (2008) Ziebart, B. D.; Maas, A.; Bagnell, J. A.; and Dey, A. K. 2008.  Maximum Entropy Inverse Reinforcement Learning.  In _Proc. AAAI_ , 1433–1438.

Supplementary Material

## Experiments - Supplementary

The gridworlds used in the main paper can be seen in Figures 1 and 2. Recall
that the third gridworld is a modification of the coffee gridworld where one
of the machines produces weak coffee, yielding a reward of $1$ when brought to
the office, and the other produces strong coffee, yielding a reward of $2$.
The original coffee world is stochastic; on each move, the agent has a $10\%$
chance of slipping orthogonal to the intended direction ($5\%$ for each
direction). The remaining two gridworlds are deterministic.

Figure 1: Recharge gridworld environment from (Vazquez-Chanlatte et al. 2017)
Figure 2: Office gridworld environment from (Toro Icarte et al. 2020)

We can analyze the inferred reward functions qualitatively by looking at
Figures 3, 4, and 5. In each, an arrow between two states with label $\ell\mid
r$ indicates that we transition from the first state to the second on input
symbol $\ell$ and receive reward $r$. We use $\vee$ to indicate multiple
symbols that cause the same transition. If we leave out $\ell$ entirely, this
indicates that all input symbols yield the same transition.

Figure 3 is the inferred reward machine for the recharge gridworld; $t$
corresponds to the drying towel, $w$ to the water, $l$ to the lava, and $r$ to
the recharging station. The agent has correctly inferred that reaching the
recharge station is rewarding while touching lava is not. Unfortunately, it
has failed to infer that recharging while wet is _not_ rewarding. The expert
rarely exhibits this behavior so it does not have much weight in the
demonstration. The learning agent is still able to perform well, however,
since there is no water on the optimal path.

Figure 4 is the inferred reward machine for the basic coffee gridworld; $c$
corresponds to the coffee, $o$ to the office, and $*$ to the decorations. The
learner has in fact inferred exactly the correct reward machine.

Finally, Figure 5 is the inferred reward machine for the multi coffee
gridworld. Here $c$ corresponds to the strong coffee and $k$ to the weak. This
machine is almost exactly correct. The only error is that the agent believes
that if it observes $*$ from state $y_{1}$ (in which it is holding the strong
coffee), it transitions to state $y_{2}$ (which should correspond to holding
the weak coffee). Rather, it should transition to the terminal state $y_{3}$.
This does not affect optimal behavior, since on this hypothesis, the agent
should still pick up the strong coffee and avoid stepping on any decorations.
Notably, the agent correctly infers the non-binary structure of the reward
function, learning that the strong coffee yields a reward of $2$ but the weak
coffee only $1$.

$y_{0}$start$y_{1}$$y_{2}$$t\vee w\mid 0$$l\mid 0$$r\mid 1$$0$$0$

Figure 3: Inferred reward machine for recharge gridworld.

$y_{0}$start$y_{1}$$y_{2}$$o\mid 0$$c\mid 0$$*\mid 0$$c\mid 0$$*\mid 0$$o\mid
1$$0$

Figure 4: Inferred reward machine for basic coffee gridworld.

$y_{0}$start$y_{1}$$y_{2}$$y_{3}$$o\mid 0$$c\mid 0$$k\mid 0$$*\mid 0$$c\vee
k\mid 0$$*\mid 0$$o\mid 2$$c\vee k\mid 0$$*\mid 0$$o\mid 1$$0$

Figure 5: Inferred reward machine for multi coffee gridworld.

## Pseudocode

Algorithm 1 Non-Markovian BIRL

Input: expert demonstration $\mathcal{D}$
Output: estimate of maximum reward machine $R_{max}$

1: $p:=p_{0}$

2: $T:=T_{0}$

3: Choose a valid RM $R_{1}$ uniformly at random

4: $R_{max}:=R_{1}$

5: Run policy iteration on $M\times R_{1}$ to compute its optimal Q-values

6: for $1\leq 1\leq N$ do

7: Choose valid neighboring RM $R_{2}$ by independently updating the
reward/target state of each transition of $R_{1}$ with probability $p$ to a
new uniformly random value

8: Run policy iteration on $M\times R_{2}$ to compute its optimal Q-values

9: // Update highest scoring RM

10: if $\left(\frac{P(\mathcal{D}\mid R_{2})}{P(\mathcal{D}\mid
R_{max})}\right)^{1/T_{min}}\frac{P(R_{2})}{P(R_{max})}>1$ then

11: $R_{max}:=R_{2}$

12: end if

13: With probability $\min\left(1,\left(\frac{P(\mathcal{D}\mid
R_{2})}{P(\mathcal{D}\mid
R_{1})}\right)^{1/T}\frac{P(R_{2})}{P(R_{1})}\right)$, set $R_{1}:=R_{2}$

14: // Update parameters every $k$ steps

15: if $i\mod k=0$ then

16: $T:=\max(T\cdot\beta_{T},T_{min})$

17: $p:=\max(p\cdot\beta_{p},p_{min})$

18: end if

19: end for

20: return $R_{max}$

For the sake of clarity, we have included pseudocode for our algorithm in
Algorithm 1.