Learning Reward Machines for Partially
Observable Reinforcement Learning
Rodrigo Toro Icarte
University of Toronto
Vector InstituteEthan Waldie
University of TorontoToryn Q. Klassen
University of Toronto
Vector Institute
Richard Valenzano
Element AIMargarita P. Castro
University of TorontoSheila A. McIlraith
University of Toronto
Vector Institute
Abstract
Reward Machines (RMs) provide a structured, automata-based representation of a
reward function that enables a Reinforcement Learning (RL) agent to decompose
an RL problem into structured subproblems that can be efﬁciently learned via
off-policy learning. Here we show that RMs can be learned from experience,
instead of being speciﬁed by the user, and that the resulting problem decomposition
can be used to effectively solve partially observable RL problems. We pose the
task of learning RMs as a discrete optimization problem where the objective is
to ﬁnd an RM that decomposes the problem into a set of subproblems such that
the combination of their optimal memoryless policies is an optimal policy for the
original problem. We show the effectiveness of this approach on three partially
observable domains, where it signiﬁcantly outperforms A3C, PPO, and ACER, and
discuss its advantages, limitations, and broader potential.1
1 Introduction
The use of neural networks for function approximation has led to many recent advances in Rein-
forcement Learning (RL) . Such deep RL methods have allowed agents to learn effective policies in
many complex environment including board games [ 30], video games [ 23], and robotic systems [ 2].
However, RL methods (including deep RL methods) often struggle when the environment is partially
observable . This is because agents in such environments usually require some form of memory to
learn optimal behaviour [ 31]. Recent approaches for giving memory to an RL agent either rely on
recurrent neural networks [24, 15, 37, 29] or memory-augmented neural networks [25, 18].
In this work, we show that Reward Machines (RMs) are another useful tool for providing memory in
a partially observable environment. RMs were originally conceived to provide a structured, automata-
based representation of a reward function [ 33,4,14,39]. Exposed structure can be exploited by the
Q-Learning for Reward Machines (QRM) algorithm [ 33], which simultaneously learns a separate
policy for each state in the RM. QRM has been shown to outperform standard and hierarchical deep
RL over a variety of discrete and continuous domains. However, QRM was only deﬁned for fully
observable environments. Furthermore, the RMs were handcrafted.
In this paper, we propose a method for learning an RM directly from experience in a partially
observable environment, in a manner that allows the RM to serve as memory for an RL algorithm.
Correspondence to: Rodrigo Toro Icarte < rntoro@cs.toronto.edu >.
1Our code is available at https://bitbucket.org/RToroIcarte/lrm .
33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada.
A requirement is that the RM learning method be given a ﬁnite set of detectors for properties that
serve as the vocabulary for the RM. We characterize an objective for RM learning that allows us to
formulate the task as a discrete optimization problem and propose an efﬁcient local search approach
to solve it. By simultaneously learning an RM and a policy for the environment, we are able to
signiﬁcantly outperform several deep RL baselines that use recurrent neural networks as memory in
three partially observable domains. We also extend QRM to the case of partial observability where
we see further gains when combined with our RM learning method.
2 Preliminaries
RL agents learn policies from experience. When the problem is fully-observable, the underlying
environment model is typically assumed to be a Markov Decision Process (MDP) . An MDP is a tuple
M=hS;A;r;p;
i, whereSis a ﬁnite set of states ,Ais a ﬁnite set of actions ,r:SA!Ris the
reward function ,p(s;a;s0)is the transition probability distribution , and
is the discount factor . The
agent starts not knowing what rorpare. At every time step t, the agent observes the current state
st2Sand executes an action at2Afollowing a policy (atjst). As a result, the state stchanges
tost+1p(st+1jst;at)and the agent receives a reward signal r(st;at). The goal is to learn the
optimal policy, which maximizes the future expected discounted reward for every state in S[32].
Q-learning [38] is a well-known RL algorithm that uses samples of experience of the form
(st;at;rt;st+1)to estimate the optimal q-function q(s;a). Here,q(s;a)is the expected return of
selecting action ain statesand following an optimal policy . Deep RL methods like DQN [ 23]
and DDQN [ 35] represent the q-function as ~q(s;a), where ~qis a neural network whose inputs are
features of the state and action, and whose weights are updated using stochastic gradient descent.
In partially observable problems, the underlying environment model is typically assumed to
be a Partially Observable Markov Decision Process (POMDP) . A POMDP is a tuple PO=
hS;O;A;r;p;!;
i, whereS,A,r,p, and
are deﬁned as in an MDP, Ois a ﬁnite set of ob-
servations , and!(s;o)is the observation probability distribution . At every time step t, the agent is
in exactly one state st2S, executes an action at2A, receives reward rt=r(st;at), and moves to
statest+1according to p(st;at;st+1). However, the agent does not observe st+1, but only receives
an observation ot+12O. This observation provides the agent a clue about what the state st+12Sis
via!. In particular, !(st+1;ot+1)is the probability of observing ot+1from statest+1[5].
RL methods cannot be immediately applied to POMDPs because the transition probabilities and
reward function are not necessarily Markovian w.r.t. O(though by deﬁnition they are w.r.t. S). As
such, optimal policies may need to consider the complete history o0;a0;:::;at1;otof observations
and actions when selecting the next action. Several partially observable RL methods use a recurrent
neural network to compactly represent the history, and then use a policy gradient method to train it.
However, when we do have access to a full POMDP model PO, then the history can be summarized
into a belief state . A belief state is a probability distribution bt:S![0;1]overS, such thatbt(s)is
the probability that the agent is in state s2Sgiven the history up to time t. The initial belief state is
computed using the initial observation o0:b0(s)/!(s;o0)for alls2S. The belief state bt+1is
then determined from the previous belief state bt, the executed action at, and the resulting observation
ot+1asbt+1(s0)/!(s0;ot+1)P
s2Sp(s;at;s0)bt(s)for alls02S. Since the state transitions and
reward function are Markovian w.r.t. bt, the set of all belief states Bcan be used to construct the
belief MDPMB. Optimal policies for MBare also optimal for the POMDP [5].
3 Reward Machines for Partially Observable Environments
In this section, we deﬁne RMs for the case of partial observability. We use the following problem as
a running example to help explain various concepts.
Example 3.1 (The cookie domain) .The cookie domain , shown in Figure 1a, has three rooms
connected by a hallway. The agent (purple triangle) can move in the four cardinal directions. There
is a button in the yellow room that, when pressed, causes a cookie to randomly appear in the red or
blue room. The agent receives a reward of +1for reaching (and thus eating) the cookie and may
then go and press the button again. Pressing the button before reaching a cookie will move it to a
random location. There is no cookie at the beginning of the episode. This is a partially observable
environment since the agent can only see what it is in the room that it currently occupies.
2
(a) Cookie domain.|
||
(b) Symbol domain.K
ὑ1 ὑ1
(c) 2-keys domain.
Figure 1: Partially observable environments where the agent can only see what is in the current room.
RMs are ﬁnite state machines that are used to encode a reward function [ 33]. In the case of partial
observability, they are deﬁned over a set of propositional symbols Pthat correspond to a set of
high-level features that the agent can detect using a labelling function L:O;A;O!2P
where (for any set X)X;,X[f;g .Lassigns truth values to symbols in Pgiven an environment
experiencee= (o;a;o0)whereo0is the observation seen after executing action awhen observing o.
We useL(;;;;o)to assign truth values to the initial observation. We call a truth value assignment
ofPanabstract observation because it provides a high-level view of the low-level environment
observations via the labelling function L. A formal deﬁnition of an RM follows:
Deﬁnition 3.1 (reward machine) .Given a set of propositional symbols P, a Reward Machine is
a tupleRP=hU;u 0;u;riwhereUis a ﬁnite set of states, u02Uis an initial state, uis the
state-transition function, u:U2P!U, andris the reward-transition function, r:U2P!R.
RMs decompose problems into a set of high-level states Uand deﬁne transitions using if-like
conditions deﬁned by u. These conditions are over a set of binary properties Pthat the agent can
detect using L. For example, in the cookie domain, P=f;;;;;;g. These properties
are true (i.e., part of an experience label according to L) in the following situations: ;;, or is
true if the agent ends the experience in a room of that color; is true if the agent ends the experience
in the same room as a cookie; is true if the agent pushed the button with its last action; and is
true if the agent ate a cookie with its last action (by moving onto the space where the cookie was).
Figure 2 shows three possible RMs for the cookie domain. They all deﬁne the same reward signal (1
for eating a cookie and 0 otherwise) but differ in their states and transitions. As a result, they differ
with respect to the amount of information about the current domain state that can be inferred from the
current RM state, as we will see below.
Each RM starts in the initial state u0. Edge labels in the ﬁgures provide a visual representation of the
functionsuandr. For example, label h;1ibetween state u2andu0in Figure 2b represents
u(u2;f;g) =u0andr(u2;f;g) = 1 . Intuitively, this means that if the RM is in state u2
and the agent’s experience ended in room immediately after eating the cookie , then the agent
will receive a reward of 1and the RM will transition to u0. Notice that any properties not listed
in the label are false (e.g. must be false to take the transition labelled h;1i). We also use
multiple labels separated by a semicolon (e.g., “ h;0i;h;0i”) to describe different conditions
for transitioning between the RM states, each with their own associated reward. The label ho/w;ri
(“o/w” for “otherwise”) on an edge from uitoujmeans that that transition will be made (and reward
rreceived) if none of the other transitions from uican be taken.
Let us illustrate the behaviour of an RM using the one shown in Figure 2c. The RM will stay in
u0until the agent presses the button (causing a cookie to appear), whereupon the RM moves to u1.
Fromu1the RM may move to u2oru3depending on whether the agent ﬁnds a cookie when it enters
another room. It is also possible to associate meaning with being in RM states: u0means that there is
no cookie available, u1means that there is a cookie in some room (either blue or red), etc.
When learning a policy for a given RM, one simple technique is to learn a policy (o;u)that considers
the current observation o2Oand the current RM state u2U. Interestingly, a partially observable
problem might be non-Markovian over O, but Markovian over OUfor some RMRP. This is the
case for the cookie domain with the RM from Figure 2c, for example.
Q-Learning for RMs (QRM) is another way to learn a policy by exploiting a given RM [ 33]. QRM
learns one q-function ~qu(i.e., policy) per RM state u2U. Then, given any sample experience,
the RM can be used to emulate how much reward would have been received had the RM been in
3
u0h;1i;
h;1i;
ho/w;0i
(a) Naive RM.u0u1 u2
ho/w;0iho/w;0iho/w;0i
h;0ih;0i;
h;0i
h;1i h;1i
(b) “Optimal” RM.u0u1 u2 u3
ho/w;0iho/w;0iho/w;0i ho/w;0i
h;0ih;0i;
h;0ih;0i;
h;0i
h;1i h;1ih;0i h;0i
(c) Perfect RM.
Figure 2: Three possible Reward Machines for the Cookie domain.
any one of its states. Formally, experience e= (o;a;o0)can be transformed into a valid experience
(ho;ui;a;ho0;u0i;r)used for updating ~qufor eachu2U, whereu0=u(u;L(e))andr=
r(u;L(e)). Hence, any off-policy learning method can take advantage of these “synthetically"
generated experiences to update all subpolicies simultaneously.
When tabular q-learning is used, QRM is guaranteed to converge to an optimal policy on fully-
observable problems [ 33]. However, in a partially observable environment, an experience emight
be more or less likely depending on the RM state that the agent was in when the experience was
collected. For example, experience emight be possible in one RM state uibut not in RM state uj.
Thus, updating the policy for ujusingeas QRM does, would introduce an unwanted bias to ~quj. We
will discuss how to (partially) address this problem in §5.
4 Learning Reward Machines from Traces
Our overall idea is to search for an RM that can be used as external memory by an agent for a given
task. As input, our method will only take a set of high-level propositional symbols P, and a labelling
functionLthat can detect them. Then, the key question is what properties should such an RM have.
Three proposals naturally emerge from the literature. The ﬁrst comes from the work on learning
Finite State Machines (FSMs) [ 3,40,10], which suggests learning the smallest RM that correctly
mimics the external reward signal given by the environment, as in Giantamidis and Tripakis’ method
for learning Moore Machines [ 10]. Unfortunately, such approaches would learn RMs of limited
utility, like the one in Figure 2a. This naive RM correctly predicts reward in the cookie domain (i.e.,
+1for eating a cookie , zero otherwise) but provides no memory in support of solving the task.
The second proposal comes from the literature on learning Finite State Controllers (FSC) [ 22] and on
model-free RL methods [ 32]. This work suggests looking for the RM whose optimal policy receives
the most reward. For instance, the RM from Figure 2b is “optimal” in this sense. It decomposes the
problem into three states. The optimal policy for u0goes directly to press the button, the optimal
policy foru1goes to the blue room and eats the cookie if present, and the optimal policy for u2goes
to the red room and eats the cookie. Together, these three policies give rise to an optimal policy for
the complete problem. This is a desirable property for RMs, but requires computing optimal policies
in order to compare the relative quality of RMs, which seems prohibitively expensive. However, we
believe that ﬁnding ways to efﬁciently learn “optimal” RMs is a promising future work direction.
Finally, the third proposal comes from the literature on Predictive State Representations (PSR)
[20], Deterministic Markov Models (DMMs) [ 21], and model-based RL [ 16]. These works suggest
learning the RM that remembers sufﬁcient information about the history to make accurate Markovian
predictions about the next observation. For instance, the cookie domain RM shown in Figure 2c
isperfect w.r.t. this criterion. Intuitively, every transition in the cookie environment is already
Markovian except for transitioning from one room to another. Depending on different factors, when
entering to the red room there could be a cookie there (or not). The perfect RM is able to encode
such information using 4 states: when at u0the agent knows that there is no cookie, at u1the agent
knows that there is a cookie in the blue or the red room, at u2the agent knows that there is a cookie
4
in the red room, and at u3the agent knows that there is a cookie in the blue room. Since keeping
track of more information will not result in better predictions, this RM is perfect . Below, we develop
a theory about perfect RMs and describe an approach to learn them.
4.1 Perfect Reward Machines: Formal Deﬁnition and Properties
The key insight behind perfect RMs is to use their states Uand transitions uto keep track of relevant
past information such that the partially observable environment PObecomes Markovian w.r.t. OU.
Deﬁnition 4.1 (perfect reward machine) .An RMRP=hU;u 0;u;riis considered perfect for a
POMDPPO=hS;O;A;r;p;!;
iwith respect to a labelling function Lif and only if for every
traceo0;a0;:::;ot;atgenerated by any policy over PO, the following holds:
Pr(ot+1;rtjo0;a0;:::;ot;at) = Pr(ot+1;rtjot;xt;at) (1)
wherex0=u0andxt=u(xt1;L(ot1;at1;ot)).
Two interesting properties follow from Deﬁnition 4.1. First, if the set of belief states Bfor the
POMDPPOis ﬁnite, then there exists a perfect RM for POwith respect to some L. Second, the
optimal policies for perfect RMs are also optimal for the POMDP (see supplementary material §12).
Theorem 4.1. Given any POMDP POwith a ﬁnite reachable belief space, there will always exists at
least one perfect RM for POwith respect to some labelling function L.
Theorem 4.2. LetRPbe a perfect RM for a POMDP POw.r.t. a labelling function L, then any
optimal policy forRPw.r.t. the environmental reward is also optimal for PO.
4.2 Perfect Reward Machines: How to Learn Them
We now consider the problem of learning a perfect RM from traces, assuming one exists w.r.t. the
given labelling function L. Recall that a perfect RM transforms the original problem into a Markovian
problem over OU. Hence, we should prefer RMs that accurately predict the next observation
o0and immediate reward rfrom the current observation o, RM stateu, and action a. This might
be achieved by collecting a training set of traces from the environment, ﬁtting a predictive model
forPr(o0;rjo;u;a ), and picking the RM that makes better predictions. However, this can be very
expensive, especially considering that the observations might be images.
Instead, we propose an alternative that focuses on a necessary condition for a perfect RM: the RM
must predict what is possible andimpossible in the environment at the abstract level. For example,
it is impossible to be at u3in the RM from Figure 2c and make the abstract observation f;g,
because the RM reaches u3only if the cookie was seen in the blue room or not to be in the red room.
This idea is formalized in the optimization model LRM. LetT=fT0;:::;Tngbe a set of traces,
where each traceTiis a sequence of observations, actions, and rewards:
Ti= (oi;0;ai;0;ri;0;:::;ai;ti1;ri;ti1;oi;ti): (2)
We now look for an RM hU;u 0;u;rithat can be used to predict L(ei;t+1)fromL(ei;t)and the
current RM state xi;t, whereei;t+1is the experience (oi;t;ai;t;oi;t+1)andei;0is(;;;;oi;0)by
deﬁnition. The model parameters are the set of traces T, the set of propositional symbols P, the
labelling function L, and a maximum number of states in the RM umax. The model also uses the sets
I=f0:::ngandTi=f0:::ti1g, whereIcontains the index of the traces and Titheir time
steps. The model has two auxiliary variables xi;tandNu;l. Variablexi;t2Urepresents the state of
the RM after observing trace Tiup to timet. VariableNu;l22Pis the set of all the next abstract
observations seen from the RM state uand the abstract observations lat some point inT. In other
words,l02Nu;liffu=xi;t,l=L(ei;t), andl0=L(ei;t+1)for some traceTiand timet.
minimize
hU;u0;u;riX
i2IX
t2Tilog(jNxi;t;L(ei;t)j) (LRM)
s:t:hU;u 0;u;ri2RP (3)
jUjumax (4)
xi;t2U 8i2I;t2Ti[ftig (5)
xi;0=u0 8i2I (6)
xi;t+1=u(xi;t;L(ei;t+1)) 8i2I;t2Ti (7)
5
Nu;l22P8u2U;l22P(8)
L(ei;t+1)2Nxi;t;L(ei;t) 8i2I;t2Ti (9)
Constraints (3)and(4)ensure that we ﬁnd a well-formed RM over Pwith at most umaxstates.
Constraint (5),(6), and (7)ensure thatxi;tis equal to the current state of the RM, starting from u0
and following u. Constraint (8)and(9)ensure that the sets Nu;lcontain every L(ei;t+1)that has
been seen right after landuinT. The objective function comes from maximizing the log-likelihood
for predicting L(ei;t+1)using a uniform distribution over all the possible options given by Nu;l.
A key property of this formulation is that any perfect RM is optimal w.r.t. the objective function in
LRMwhen the number of traces tends to inﬁnity (see supplementary material §12):
Theorem 4.3. When the set of training traces (and their lengths) tends to inﬁnity and is collected by
a policy such that (ajo)>for allo2Oanda2A, any perfect RM with respect to Land at most
umaxstates will be an optimal solution to the formulation LRM.
Finally, note that the deﬁnition of a perfect RM does not impose conditions over the rewards associated
with the RM (i.e., r). This is why ris a free variable in the model LRM. However, we still expect r
to model the external reward signals given by the environment. To do so, we estimate r(u;l)using
its empirical expectation over T(as commonly done when constructing belief MDPs [5]).
4.3 Searching for a Perfect Reward Machine Using Tabu Search
We now describe the speciﬁc optimization technique used to solve LRM. We experimented with many
discrete optimization approaches—including mixed integer programming [ 6], Benders decomposi-
tion [ 8], evolutionary algorithms [ 17], among others—and found local search algorithms [ 1] to be
the most effective at ﬁnding high quality RMs given short time limits. In particular, we use Tabu
search [ 11], a simple and versatile local search procedure with convergence guarantees and many
successful applications in the literature [ 36]. We also include our unsuccessful mixed integer linear
programming model for LRMin the supplementary material §10.
In the context of our work, Tabu search starts from a random RM and, on each iteration it evaluates all
“neighbouring” RMs. We deﬁne the neighbourhood of an RM as the set of RMs that differ by exactly
one transition (i.e., removing/adding a transition, or changing its value) and evaluate RMs using the
objective function of LRM. When all neighbouring RMs are evaluated, the algorithm chooses the one
with lowest values and sets it as the current RM. To avoid local minima, Tabu search maintains a
Tabu list of all the RMs that were previously used as the current RM. Then, RMs in the Tabu list are
pruned when examining the neighbourhood of the current RM.
5 Simultaneously Learning a Reward Machine and a Policy
We now describe our overall approach to simultaneously ﬁnding an RM and exploiting that RM to
learn a policy. The complete pseudo-code can be found in the supplementary material (Algorithm 1).
Our approach starts by collecting a training set of traces Tgenerated by a random policy during tw
“warmup” steps. This set of traces is used to ﬁnd an initial RM Rusing Tabu search. The algorithm
then initializes policy , sets the RM state to the initial state u0, and sets the current label lto the
initial abstract observation L(;;;;o). The standard RL learning loop is then followed: an action
ais selected following (o;u)whereuis the current RM state, and the agent receives the next
observation o0and the immediate reward r. The RM state is then updated to u0=u(u;L(o;a;o0))
and the last experience (ho;ui;a;r;ho0;u0i)is used by any RL method of choice to update . Note
that in an episodic task, the environment and RM are reset whenever a terminal state is reached.
If on any step, there is evidence that the current RM might not be the best one, our approach will
attempt to ﬁnd a new one. Recall that the RM Rwas selected using the cardinality of its prediction
setsN(LRM). Hence, if the current abstract observation l0is not inNu;l, adding the current trace to T
will increase the size of Nu;lforR. As such, the cost of Rwill increase and it may no longer be the
best RM. Thus, if l062Nu;l, we add the current trace to Tand search for a new RM. Recall that we
use Tabu search, though any discrete optimization method could be applied. Our method only uses
the new RM if its cost is lower than R’s. If the RM is updated, a new policy is learned from scratch.
Given the current RM, we can use any RL algorithm to learn a policy (o;u), by treating the
combination of oanduas the current state. If the RM is perfect, then the optimal policy (o;u)will
6
Cookie domain Symbol domain 2-keys domains
0 1·1062·1063·106050100150200
Training stepsReward
0 1·1062·1060200400
Training stepsReward
0 2·1064·106050100150
Training stepsReward
Legend: DDQN A3C PPO ACER LRM + DDQN LRM + DQRM Optimal
Figure 3: Total reward collected every 10;000training steps.
also be optimal for the original POMDP (as stated in Theorem 4.2). However, to exploit the problem
structure exposed by the RM, we can use the QRM algorithm.
As explained in §3, standard QRM under partial observability can introduce a bias because an
experiencee= (o;a;o0)might be more or less likely depending on the RM state that the agent was
in when the experience was collected. We partially address this issue by updating ~quusing (o;a;o0)
if and only if L(o;a;o0)2Nu;l, wherelwas the current abstract observation that generated the
experience (o;a;o0). Hence, we do not transfer experiences from uitoujif the current RM does not
believe that (o;a;o0)is possible in uj. For example, consider the cookie domain and the perfect RM
from Figure 2c. If some experience consists of entering to the red room and seeing a cookie, then
this experience will not be used by states u0andu3as it is impossible to observe a cookie at the red
room from those states. Note that adding this rule may work in many cases, but it will not address
the problem in all environments (more discussion in §7). We consider addressing this problem as an
interesting area for future work.
6 Experimental Evaluation
We tested our approach on three partially observable grid domains (Figure 1). The agent can move
in the four cardinal directions and can only see what is in the current room. These are stochastic
domains where the outcome of an action randomly changes with a 5% probability.
The ﬁrst environment is the cookie domain (Figure 1a) described in §3. Each episode is 5;000steps
long, during which the agent should attempt to get as many cookies as possible.
The second environment is the symbol domain (Figure 1b). It has three symbols |,, andin the
red and blue rooms. One symbol from f|;;gand possibly a right or left arrow are randomly
placed at the yellow room. Intuitively, that symbol and arrow tell the agent where to go, e.g., |and
!tell the agent to go to |in the east room. If there is no arrow, the agent can go to the target symbol
in either room. An episode ends when the agent reaches any symbol in the red or blue room, at which
point it receives a reward of +1if it reached the correct symbol and 1otherwise.
The third environment is the 2-keys domain (Figure 1c). The agent receives a reward of +1when
it reaches the coffee (in the yellow room). To do so, it must open the two doors (shown in brown).
Each door requires a different key to open it, and the agent can only carry one key at a time. Initially,
the two keys are randomly located in either the blue room, the red room, or split between them.
We tested two versions of our Learned Reward Machine (LRM) approach: LRM+DDQN and
LRM+DQRM. Both learn an RM from experience as described in §4.2, but LRM+DDQN learns
a policy using DDQN [ 35] while LRM+DQRM uses the modiﬁed version of QRM described in
§5. In all domains, we used umax= 10 ,tw= 200;000, an epsilon greedy policy with = 0:1, and
a discount factor
= 0:9. The size of the Tabu list and the number of steps that the Tabu search
performs before returning the best RM found is 100. We compared against 4 baselines: DDQN [ 35],
A3C [ 24], ACER [ 37], and PPO [ 29] using the OpenAI baseline implementations [ 12]. DDQN uses
the concatenation of the last 10 observations as input which gives DDQN a limited memory to better
handle the domains. A3C, ACER, and PPO use an LSTM to summarize the history. Note that the
7
output of the labelling function was also given to the baselines. Details on the hyperparameters and
networks can be found in the supplementary material §13.
Figure 3 shows the total cumulative rewards that each approach gets every 10;000training steps and
compares it to the optimal policy. For the LRM algorithms, the ﬁgure shows the median performance
over 30 runs per domain, and percentile 25 to 75 in the shadowed area. For the DDQN baseline, we
show the maximum performance seen for each time period over 5 runs per problem. Similarly, we
also show the maximum performance over the 30 runs of A3C, ACER, and PPO per period. All the
baselines outperformed a random policy, but none make much progress on any of the domains.
Furthermore, LRM approaches largely outperform all the baselines, reaching close-to-optimal policies
in the cookie and symbol domain. We also note that LRM+DQRM learns faster than LRM+DDQN,
but is more unstable. In particular, LRM+DQRM converged to a considerably better policy than
LRM+DDQN in the 2-keys domain. We believe this is due to QRM’s experience sharing mechanism
that allows for propagating sparse reward backwards faster (see supplementary material §13.3).
A key factor in the strong performance of the LRM approaches is that Tabu search ﬁnds high-quality
RMs in less than 100 local search steps (Figure 5, supplementary material). In fact, our results show
that Tabu search ﬁnds perfect RMs in most runs, in particular when tested over the symbol domain.
7 Discussion, Limitations, and Broader Potential
u0u1h;1i;
ho/w;0i
h;1i;
ho/w;0ih;0ih;0i
Figure 4: The gravity domain.Solving partially observable RL problems is challenging
and LRM was able to solve three problems that were con-
ceptually simple but presented a major challenge to A3C,
ACER, and PPO with LSTM-based memories. A key idea
behind these results was to optimize over a necessary con-
dition for perfect RMs. This objective favors RMs that
are able to predict possible and impossible future observa-
tions at the abstract level given by the labelling function
L. In this section, we discuss the advantages and current
limitations of such an approach.
We begin by considering the performance of Tabu search
in our domains. Given a training set composed of one
million transitions, a simple Python implementation of
Tabu search takes less than 2.5 minutes to learn an RM
across all our environments, when using 62 workers on a Threadripper 2990WX processor. Note that
Tabu search’s main bottleneck is evaluating the neighbourhood around the current RM solution. As
the size of the neighbourhood depends on the size of the set of propositional symbols P, exhaustively
evaluating the neighbourhood may sometimes become impractical. To handle such problem, it will
be necessary to import ideas from the Large Neighborhood Search literature [27].
Regarding limitations, learning the RM at the abstract level is efﬁcient but requires ignoring (possibly
relevant) low-level information. For instance, Figure 4 shows an adversarial example for LRM.
The agent receives reward for eating the cookie ( ). There is an external force pulling the agent
down—i.e., the outcome of the “ move-up ” action is actually a downward movement with high
probability. There is a button ( ) that the agent can press to turn off (or back on) the external force.
Hence, the optimal policy is to press the button and then eat the cookie. Given P=f;g, a perfect
RM for this environment is fairly simple (see Figure 4) but LRM might not ﬁnd it. The reason is that
pressing the button changes the low-level probabilities in the environment but does not change what
is possible or impossible at the abstract level. In other words, while the LRM objective optimizes
over necessary conditions for ﬁnding a perfect RM, those conditions are not sufﬁcient to ensure that
an optimal solution will be a perfect RM. In addition, if a perfect RM is found, our heuristic approach
to share experiences in QRM would not work as intended because the experiences collected when the
force is on (at u0) would be used to update the policy for the case where the force is off (at u1).
Other current limitations include that it is unclear how to handle noise over the high-level detectors L
and how to transfer learning from previously learned policies when a new RM is learned. Finally,
deﬁning a set of proper high-level detectors for a given environment might be a challenge to deploying
LRM. Hence, looking for ways to automate that step is an important direction for future work.
8
8 Related Work
State-of-the-art approaches to partially observable RL use Recurrent Neural Networks (RNNs) as
memory in combination with policy gradient [ 24,37,29,15], or use external neural-based memories
[25,18,13]. Other approaches include extensions to Model-Based Bayesian RL to work under partial
observability [ 28,7,9] and to provide a small binary memory to the agent and a special set of actions
to modify it [ 26]. While our experiments highlight the merits of our approach w.r.t. RNN-based
approaches, we rely on ideas that are largely orthogonal. As such, we believe there is signiﬁcant
potential in mixing these approaches to get the beneﬁt of memory at both the high- and the low-level.
The effectiveness of automata-based memory has long been recognized in the POMDP literature [ 5],
where the objective is to ﬁnd policies given a complete speciﬁcation of the environment. The idea
is to encode policies using Finite State Controllers (FSCs) which are FSMs where the transitions
are deﬁned in terms of low-level observations from the environment and each state in the FSM is
associated with one primitive action. When interacting with the environment, the agent always selects
the action associated with the current state in the controller. Meuleau et al. [22] adapted this idea to
work in the RL setting by exploiting policy gradient to learn policies encoded as FSCs. RMs can be
considered as a generalization of FSC as they allow for transitions using conditions over high-level
events and associate complete policies (instead of just one primitive action) to each state. This allows
our approach to easily leverage existing deep RL methods to learn policies from low-level inputs,
such as images—which is not achievable by Meuleau et al. [22]. That said, further investigating using
ideas for learning FSMs [3, 40, 10] in learning RMs is a promising direction for future work.
Our approach to learn RMs is greatly inﬂuenced by Predictive State Representations (PSRs) [ 20].
The idea behind PSRs is to ﬁnd a set of core tests (i.e., sequences of actions and observations) such
that if the agent can predict the probabilities of these occurring, given any history H, then those
probabilities can be used to compute the probability of any other test given H. The insight is that
state representations that are good for predicting the next observation are good for solving partially
observable environments. We adapted this idea to the context of RM learning as discussed in §4.
While our work was under review, two interesting papers were submitted to arXiv. The ﬁrst paper, by
Xu et al. [ 39], proposes a polynomial time algorithm to learn reward machines in fully observable
domains. Their goal is to learn the smallest reward machine that is consistent with the reward
function—which makes sense for fully observable domains, but would have limited utility under
partial observability (as discussed in §4). The second paper, by Zhang et al. [ 41], proposes to learn a
discrete PSR representation of the environment directly from low-level observations and then plan
over such representation using tabular Q-learning. This is a promising research direction, with some
clear synergies with LRM.
9 Concluding Remarks
We have presented a method for learning (perfect) Reward Machines in partially observable envi-
ronments and demonstrated the effectiveness of these learned RMs in tackling partially observable
RL problems that are unsolvable by A3C, ACER and PPO. Informed by criteria from the POMDP,
FSC, and PSR literature, we proposed a set of RM properties that support tackling RL in partially
observable environments. We used these properties to formulate RM learning as a discrete optimiza-
tion problem. We experimented with several optimization methods, ﬁnding Tabu search to be the
most effective. We then combined this RM learning with policy learning for partially observable
RL problems. Our combined approach outperformed a set of strong LSTM-based approaches on
different domains.
We believe this work represents an important building block for creating RL agents that can solve
cognitively challenging partially observable tasks. Not only did our approach solve problems that
were unsolvable by A3C, ACER and PPO, but it did so in a relatively small number of training steps.
RM learning provided the agent with memory, but more importantly the combination of RM learning
and policy learning provided it with discrete reasoning capabilities that operated at a higher level of
abstraction, while leveraging deep RL’s ability to learn policies from low-level inputs. This work
leaves open many interesting questions relating to abstraction, observability, and properties of the
language over which RMs are constructed. We believe that addressing these questions, among many
others, will push the boundary of partially observable RL problems that can be solved.
9
Acknowledgments
We gratefully acknowledge funding from the Natural Sciences and Engineering Research Council of
Canada (NSERC) and Microsoft Research. The ﬁrst author also gratefully acknowledges funding
from CONICYT (Becas Chile). A preliminary version of this work was presented at RLDM [34].
References
[1]E. Aarts, E. H. Aarts, and J. K. Lenstra. Local search in combinatorial optimization . Princeton
University Press, 2003.
[2]M. Andrychowicz, B. Baker, M. Chociej, R. Jozefowicz, B. McGrew, J. Pachocki, A. Petron,
M. Plappert, G. Powell, A. Ray, et al. Learning dexterous in-hand manipulation. arXiv preprint
arXiv:1808.00177 , 2018.
[3]D. Angluin and C. H. Smith. Inductive inference: Theory and methods. ACM Computing
Surveys (CSUR) , 15(3):237–269, 1983.
[4]A. Camacho, R. Toro Icarte, T. Q. Klassen, R. Valenzano, and S. A. McIlraith. LTL and beyond:
Formal languages for reward function speciﬁcation in reinforcement learning. In Proceedings
of the 28th International Joint Conference on Artiﬁcial Intelligence (IJCAI) , pages 6065–6073,
2019.
[5]A. R. Cassandra, L. P. Kaelbling, and M. L. Littman. Acting optimally in partially observable
stochastic domains. In Proceedings of the 12th National Conference on Artiﬁcial Intelligence
(AAAI) , pages 1023–1028, 1994.
[6]M. Conforti, G. Cornuéjols, and G. Zambelli. Integer programming , volume 271. Springer,
2014.
[7]F. Doshi-Velez, D. Pfau, F. Wood, and N. Roy. Bayesian nonparametric methods for partially-
observable reinforcement learning. IEEE transactions on pattern analysis and machine intelli-
gence , 37(2):394–407, 2013.
[8]A. M. Geoffrion. Generalized Benders decomposition. Journal of optimization theory and
applications , 10(4):237–260, 1972.
[9]M. Ghavamzadeh, S. Mannor, J. Pineau, A. Tamar, et al. Bayesian reinforcement learning: A
survey. Foundations and Trends in Machine Learning , 8(5-6):359–483, 2015.
[10] G. Giantamidis and S. Tripakis. Learning Moore machines from input-output traces. In
Proceedings of the 21st International Symposium on Formal Methods (FM) , pages 291–309,
2016.
[11] F. Glover and M. Laguna. Tabu search. In Handbook of combinatorial optimization , pages
2093–2229. Springer, 1998.
[12] C. Hesse, M. Plappert, A. Radford, J. Schulman, S. Sidor, and Y . Wu. OpenAI baselines.
https://github.com/openai/baselines , 2017.
[13] C.-C. Hung, T. Lillicrap, J. Abramson, Y . Wu, M. Mirza, F. Carnevale, A. Ahuja, and
G. Wayne. Optimizing agent behavior over long time scales by transporting value. arXiv
preprint arXiv:1810.06721 , 2018.
[14] L. Illanes, X. Yan, R. Toro Icarte, and S. A. McIlraith. Symbolic planning and model-free
reinforcement learning: Training taskable agents. In Proceedings of the 4th Multi-disciplinary
Conference on Reinforcement Learning and Decision (RLDM) , pages 191–195, 2019.
[15] M. Jaderberg, V . Mnih, W. M. Czarnecki, T. Schaul, J. Z. Leibo, D. Silver, and K. Kavukcuoglu.
Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397 ,
2016.
[16] L. P. Kaelbling, M. L. Littman, and A. W. Moore. Reinforcement learning: A survey. Journal
of artiﬁcial intelligence research , 4:237–285, 1996.
[17] D. Kasenberg and M. Scheutz. Interpretable apprenticeship learning with temporal logic
speciﬁcations. In Proceedings of the 56th IEEE Annual Conference on Decision and Control
(CDC) , pages 4914–4921, 2017.
[18] A. Khan, C. Zhang, N. Atanasov, K. Karydis, V . Kumar, and D. D. Lee. Memory augmented
control networks. arXiv preprint arXiv:1709.05706 , 2017.
10
[19] D. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint
arXiv:1412.6980 , 2014.
[20] M. L. Littman, R. S. Sutton, and S. Singh. Predictive representations of state. In Proceedings
of the 15th Conference on Advances in Neural Information Processing Systems (NIPS) , pages
1555–1561, 2002.
[21] M. Mahmud. Constructing states for reinforcement learning. In Proceedings of the 27th
International Conference on Machine Learning (ICML) , pages 727–734, 2010.
[22] N. Meuleau, L. Peshkin, K.-E. Kim, and L. P. Kaelbling. Learning ﬁnite-state controllers for
partially observable environments. In Proceedings of the 15th Conference on Uncertainty in
Artiﬁcial Intelligence (UAI) , pages 427–436, 1999.
[23] V . Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves,
M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep rein-
forcement learning. Nature , 518(7540):529–533, 2015.
[24] V . Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu.
Asynchronous methods for deep reinforcement learning. In Proceedings of the 33rd Interna-
tional Conference on Machine Learning (ICML) , pages 1928–1937, 2016.
[25] J. Oh, V . Chockalingam, S. Singh, and H. Lee. Control of memory, active perception, and
action in minecraft. In Proceedings of the 33rd International Conference on Machine Learning
(ICML) , pages 2790–2799, 2016.
[26] L. Peshkin, N. Meuleau, and L. P. Kaelbling. Learning policies with external memory. In
Proceedings of the 16th International Conference on Machine Learning (ICML) , pages 307–314,
1999.
[27] D. Pisinger and S. Ropke. Large neighborhood search. In Handbook of metaheuristics , pages
399–419. Springer, 2010.
[28] P. Poupart and N. Vlassis. Model-based bayesian reinforcement learning in partially observable
domains. In Proceedings of the 10th International Symposium on Artiﬁcial Intelligence and
Mathematics (ISAIM) , pages 1–2, 2008.
[29] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization
algorithms. arXiv preprint arXiv:1707.06347 , 2017.
[30] D. Silver, J. Schrittwieser, K. Simonyan, I. Antonoglou, A. Huang, A. Guez, T. Hubert, L. Baker,
M. Lai, A. Bolton, et al. Mastering the game of Go without human knowledge. Nature , 550
(7676):354, 2017.
[31] S. P. Singh, T. Jaakkola, and M. I. Jordan. Learning without state-estimation in partially
observable Markovian decision processes. In Machine Learning Proceedings 1994 , pages
284–292. Elsevier, 1994.
[32] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018.
[33] R. Toro Icarte, T. Q. Klassen, R. Valenzano, and S. A. McIlraith. Using reward machines for
high-level task speciﬁcation and decomposition in reinforcement learning. In Proceedings of
the 35th International Conference on Machine Learning (ICML) , pages 2112–2121, 2018.
[34] R. Toro Icarte, E. Waldie, T. Q. Klassen, R. Valenzano, M. P. Castro, and S. A. McIlraith.
Searching for Markovian subproblems to address partially observable reinforcement learning. In
Proceedings of the 4th Multi-disciplinary Conference on Reinforcement Learning and Decision
(RLDM) , pages 22–26, 2019.
[35] H. Van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In
Proceedings of the 30th AAAI Conference on Artiﬁcial Intelligence (AAAI) , pages 2094–2100,
2016.
[36] S. V oß, S. Martello, I. H. Osman, and C. Roucairol. Meta-heuristics: Advances and trends in
local search paradigms for optimization . Springer Science & Business Media, 2012.
[37] Z. Wang, V . Bapst, N. Heess, V . Mnih, R. Munos, K. Kavukcuoglu, and N. de Freitas. Sample
efﬁcient actor-critic with experience replay. arXiv preprint arXiv:1611.01224 , 2016.
[38] C. J. C. H. Watkins and P. Dayan. Q-learning. Machine learning , 8(3-4):279–292, 1992.
11
[39] Z. Xu, I. Gavran, Y . Ahmad, R. Majumdar, D. Neider, U. Topcu, and B. Wu. Joint inference of
reward machines and policies for reinforcement learning. arXiv preprint arXiv:1909.05912 ,
2019.
[40] Z. Zeng, R. M. Goodman, and P. Smyth. Learning ﬁnite state machines with self-clustering
recurrent networks. Neural Computation , 5(6):976–990, 1993.
[41] A. Zhang, Z. C. Lipton, L. Pineda, K. Azizzadenesheli, A. Anandkumar, L. Itti, J. Pineau, and
T. Furlanello. Learning causal state representations of partially observable environments. arXiv
preprint arXiv:1906.10437 , 2019.
12
Learning Reward Machines for Partially
Observable Reinforcement Learning
Supplementary Material
10 Mixed Integer Linear Programming Model for LRM
In this section, we present a Mixed Integer Linear Programming model ( MILP ) for LRM. We assume
jUj=umaxand setK= 22jPj. Variablesdu;u0;l2f0;1grepresent the possible transitions in the
RM for each pair of states u;u02Uand abstract observation l22P. Formally, this means that
du;u0;l= 1iffu(u;l) =u0. Variablewi;t;u2f0;1gindicates if the agent is at state u2Uof the
RM on trace i2Iand time step t2Ti, which corresponds to the statement that wi;t;u= 1 iff
xi;t=u. Variablepl;u;l02f0;1gindicates ifl022Pis a possible next abstract observation at RM
stateuwhen observing l22P. This means that pl;u;l0= 1iffl02Nu;l. Variableyu;l;n2f0;1g
represents the cardinality of Nu;l, meaning that yu;l;n = 1 iffjNu;lj=n. Lastly, variables zi;t
represent the log-likelihood cost for trace i2Iand time step t2Ti, which can be formally stated as
zi;t= log(jNxi;t;L(ei;t)j). The full model is then as follows:
minX
i2IX
t2Tizi;t (MILP )
s:t:zi;tKX
n=1yu;l;nlog (n)(1wi;t;u)log (K)8i2I;t2Ti;u2U;l=L(ei;t)(10)
KX
n=1yu;l;n= 1 8u2U;l22P(11)
X
l022Ppl;u;l0=KX
n=1yu;l;nn 8u2U;l22P;n2f1::Kg(12)
pl;u;l0wi;t;u 8i2I;t2Ti;l=L(ei;t);l0=L(ei;t+1)(13)
X
u02Udu;u0;l= 1 8u2U;l22P(14)
X
u2Uwi;t;u= 1 8i2I;t2Ti;u2U(15)
wi;0;u0= 1 8i2I(16)
wi;t+1;u0wi;t;u+du;u0;l1 8i2I;t2Ti;u;u02U;l=L(ei;t+1)(17)
du;u0;l2f0;1g 8 u;u02U;l22P(18)
wi;t;u2f0;1g 8 i2I;t2Ti;u2U(19)
pl;u;l02f0;1g 8 u2U;l;l022P(20)
yu;l;n2f0;1g 8 u2U;l22P;n2f1::Kg(21)
zi;t0 8i2I;t2Ti(22)
Constraint (10) models the log-likelihood cost for each time step of a trace. Constraints (11) and
(12) compute the cardinality of Nu;l. Constraint (13) deﬁnes the possible predictions given a trace.
Constraint (14) enforces that for each RM state an abstract observation can lead to exactly one other
RM state. Constraint (15) enforces that at any time step of a trace the agent can be at exactly one RM
state. Constraint (16) imposes the initial RM state of a trace, and constraint (17) encodes the RM
state transitions for a trace. Finally, constraints (18)-(22) correspond to the variables’ domains.
13
11 Algorithm for Simultaneously Learning Reward Machines and a Policy
Algorithm 1 shows our overall approach to simultaneously learning an RM and exploiting that RM to
learn a policy. The algorithm inputs are the set of propositional symbols P, the labelling function L,
a maximum on the number of RM states umax, and the number of “warmup" steps tw. Our approach
starts by collecting a training set of traces Tgenerated by a random policy during twsteps (Line 2).
This set of traces is used to ﬁnd an initial RM Rusing Tabu search (Line 3). If later traces show that
Ris incorrect, our method will then ﬁnd a new RM learned using the additional traces.
Lines 4 and 5 initialize the environment and the policy , and set variables xandlto the initial RM
stateu0and initial abstract observation L(;;;;o), respectively. Lines 7–19 are the main loop of
our approach. Lines 7–10 are part of the standard RL loop: the agent executes an action aselected
following(o;x)and receives the next observation o0, the immediate reward r, and a boolean variable
done indicating if the episode has terminated. Then, the state in the RM x0is updated and the policy
is improved using the last experience (ho;xi;a;r;ho0;x0i;done). Note that when done is true, the
environment and RM are reset (Lines 17–18).
Lines 11–16 involve relearning the RM when there is evidence that the current RM might not be
the best one. Recall that the RM Rwas selected using the cardinality of its prediction sets N(see
the description of LRM). Hence, if the current abstract observation l0is not inNx;l, then adding the
current trace toTwill increase the size of Nx;lforR. As such, the cost of Rwill increase and it
may no longer be the best RM. Thus, if l062Nx;l, we add the current trace to Tand use Tabu search
to ﬁnd a new RM. Note, our method only uses the new RM if its cost is lower than that of R(Lines
14–16). However, when the RM is updated, a new policy is learned from scratch (Line 16).
Algorithm 1 Learning an RM and a Policy
1:Input:P,L,A,umax,tw
2:T  collect_traces( tw)
3:R;N learn_rm(P,L,T,umax)
4:o;x;l env_get_initial_state(), u0,L(;;;;o)
5: initialize_policy()
6:fort= 1tottraindo
7:a select_action( ,o,x)
8:o0;r;done env_execute_action( a)
9:x0;l0 u(x;L(o;a;o0));L(o;a;o0)
10: improve(,o,x,l,a,r,o0,x0,l0, done,N)
11: ifl062Nx;lthen
12:T  T[ get_current_trace()
13:R0;N relearn_rm(R,P,L,T,umax)
14: ifR6=R0then
15:R;done R0,true
16: initialize_policy()
17: end if
18: end if
19: ifdone then
20:o0;x0;l0 env_get_initial_state(), u0,L(;;;;o)
21: end if
22:o;x;l o0;x0;l0
23:end for
24:return
14
12 Theorems and Proof Sketches
Theorem 12.1. Given any POMDP POwith a ﬁnite reachable belief space, there exist a perfect RM
forPOwith respect to some labelling function L.
Proof sketch. If the reachable belief space Bis ﬁnite, we can construct an RM that keeps track of the
current belief state using one RM state per belief state and emulating their progression using u, and
one propositional symbol for every action-observation pair. Thus, the current belief state btcan be
inferred from the last observation, last action, and the current RM state. As such, the equality from
Deﬁnition 4.1 holds.
Two interesting properties follow from the deﬁnition of a perfect RM. First, if the set of belief states
Bfor the POMDPPOis ﬁnite, then there exists a perfect RM for POwith respect to some L. Second,
the optimal policies for perfect RMs are also optimal for the POMDP.
Theorem 12.2. LetRPbe a perfect RM for a POMDP POw.r.t. a labelling function L, then any
optimal policy forRPw.r.t. the environmental reward is also optimal for PO.
Proof sketch. As the next observation and immediate reward probabilities can be predicted from
OUA, an optimal policy over OUmust also be optimal over PO.
A key property of this formulation is that any perfect RM is optimal with respect to the objective
function in LRMwhen the number of traces tends to inﬁnity:
Theorem 12.3. When the set of training traces (and their lengths) tends to inﬁnity and is collected
by a policy such that (ajo)>for allo2Oanda2A, and some constant >0, then any perfect
RM with respect to Land at most umaxstates will be an optimal solution to the formulation given in
LRM.
Proof sketch. In the limit, l02Nu;lif and only if the probability of observing l0after executing an
action from the RM state uwhile observing lis non-zero. In particular, for all i2Iandt2T,
the cardinality of Nxi;t;L(ei;t)will be minimal for a perfect RM. This follows from the fact that
perfect RMs make perfect predictions for the next observation o0giveno,u, anda. Therefore, as we
minimize the sum over log(jNxi;t;L(ei;t)j), the objective value for a perfect RM must be minimal.
15
13 Experimental Evaluation
13.1 Experimental Details
For LRM+DDQN and LRM+DQRM, the neural network used has 5 fully connected layers with 64
neurons per layer. On every step, we trained the network using 32 sampled experiences from a replay
buffer of size 100,000 using the Adam optimizer [ 19] and a learning rate of 5e-5. The target networks
were updated every 100steps.
DDQN [ 35] uses the same parameters and network architecture as LRM+DDQN, but its input is the
concatenation of the last 10 observations, as commonly done by Atari playing agents. This gives
DDQN a limited memory to better handle partially observable domains. In contrast, A3C, ACER,
and PPO use an LSTM to summarize the history.
We also followed the same testing methodology that was used in their original publications. We ran
each approach at least 30 times per domain, and on every run, we randomly selected the number of
hidden neurons for the LSTM from f64;128;256;512gand a learning rate from (1e-3, 1e-5). We
also sampled fromf0;1;2gfor ACER and the clip range from (0:1;0:3)for PPO. Other parameters
were ﬁxed to their default values.
While interacting with the environment, the agents were given a “top-down" view of the world
represented as a set of binary matrices. One matrix had a 1 in the current location of the agent,
one had a 1 in only those locations that are currently observable, and the remaining matrices each
corresponded to an object in the environment and had a 1 at only those locations that were both
currently observable and contained that object (i.e., locations in other rooms are “blacked out"). The
agent also had access to features indicating if they were carrying a key, which colour room they were
in, and the current status (i.e., occurring or not occurring) of the events detected by the labelling
function.
13.2 Tabu Search
Figure 5 evaluates the quality of the RMs found by Tabu search by comparing it the perfect RM. In
each plot, a dot compares the cost of a handcrafted perfect RM with that of an RM Rthat was found
by Tabu search while running our LRM approaches, where the costs are evaluated relative to the
training set used to ﬁnd R. Being on or under the diagonal line (as in most of the points in the ﬁgure)
means that Tabu search is ﬁnding RMs whose values are at least as good as the handcrafted RM.
Hence, Tabu search is either ﬁnding perfect RMs or discovering that our training set is incomplete
and our agent will eventually ﬁll those gaps.
Cookie domain Symbol domain 2-keys domains
103.7103.8103.9104103.7103.8103.9104
Cost Perfect RMCost Learned RM
103.7103.8103.9103.7103.8103.9
Cost Perfect RMCost Learned RM
103.8104104.2103.8104104.2
Cost Perfect RMCost Learned RM
Figure 5: Cost comparison between perfect RM and RM found by Tabu search.
13.3 DDQN vs DQRM: Exploration Heatmaps and Learned Trajectories
As shown in §6 of the paper, LRM+DQRM tends to learn faster than LRM+DDQN and largely
outperforms LRM+DDQN in the 2-keys domain. Towards better understating these results, we ran
the following experiment. For the 2-keys domain (and identical random seed), we learn policies using
DDQN and DQRM over the same handcrafted perfect RM from Figure 6.
16
In the 2-keys domain, P=fὑ1;ὑ1ὑ1;;K;;;;g. These properties are true in the following
situations: ;;, or is true if the agent ends the experience in a room of that color; ὑ1is true if
the agent ends the experience in the same room as exactly one key; ὑ1ὑ1 is true if the agent ends the
experience in the same room as exactly two keys; is true if the agent is carrying a key; and Kis
true if the agent reaches the coffee with its last action.
u0u1
u2
u3u4
u5u6 ho/w;0iho/w;0i
ho/w;0i
ho/w;0iho/w;0i
ho/w;0ihK;1i;
ho/w;0ihὑ1ὑ1;0i;
h;0i
hὑ1;0i;hὑ1;0i
h;0i;
hὑ1ὑ1;0ihὑ1;0i
hὑ1;0ih;0i
h;0ih;0i
h;0i
Figure 6: A perfect reward machine for the 2-keys domain.
Figures 7 and 8 show the exploration heatmaps of DDQN and DQRM agents during the ﬁrst 500;000
training steps (red areas represent places where the agent spent more time). For both approaches,
the ﬁrst 50;000steps are random. After that, learning begins and both agents follow an -greedy
policy with= 0:1. The most interesting difference between DDQN and DQRM is between training
steps 50;000and150;000(second and third maps from the left in Figure 7). They show that DQRM
spends less time in the coffee room and the hallway than DDQN. This could be explained by the two
main differences between these RL methods, as detailed below.
DDQN DDQN DDQN DDQN DDQN
0–50K steps 50–100K steps 100–150K steps 150–200K steps 200–250K steps
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
DQRM DQRM DQRM DQRM DQRM
0–50K steps 50–100K steps 100–150K steps 150–200K steps 200–250K steps
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
Figure 7: Exploration heatmaps for DDQN and DQRM over the 2-keys domain given a perfect RM.
First, DQRM decomposes the main policy into one q-network per state in the RM. Hence, each
q-network has to learn a relatively simpler policy than DDQN. In fact, the q-network for the last
17
RM stateu6is fairly simple: when the agent is carrying a key and there is only one closed door, go
and get the coffee. As expected, DQRM seems to learn an accurate estimation of that q-function
using very limited interactions within the coffee room. In contrast, DDQN uses one big q-network to
model the complete policy. Receiving reward by getting the coffee pushes the q-network estimation
to believe that a high reward can be collected from the coffee room (even if the doors are closed).
Hence, the agent spent a considerable amount of time hitting the second door without having a key.
Second, DQRM shares experience over all the q-networks. This allows for using the experiences
collected while being at early stages of the task (e.g., states u0,u1,u2, andu3) to update policies
for later stages (e.g., states u4,u5, andu6). In particular, all the experience needed to learn to
navigate from one room to another is shared. Therefore, while DDQN would depend on its network
to avoid relearning how to navigate between rooms when the RM state changes, DQRM enforces
such transfer by sharing experiences whenever it is appropriate. This might explain why DDQN
spends considerably more time in the hallway than DQRM in heatmaps 50-100K and 100-150K.
DDQN DDQN DDQN DDQN DDQN
250–300K steps 300–350K steps 350–400K steps 400–450K steps 450–500K steps
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
DQRM DQRM DQRM DQRM DQRM
250–300K steps 300–350K steps 350–400K steps 400–450K steps 450–500K steps
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
K
ὑ1ὑ1
Figure 8: Exploration heatmaps for DDQN and DQRM over the 2-keys domain given a perfect RM.
Note that the exploratory trend from 150-500K follows a clear pattern. On one hand, the DQRM
agent seems to have a good idea of how to solve the task and, therefore, it spends most of its time on
the hallway (solving the tasks requires passing through the hallway at least 4 times). On the other
hand, the DDQN agent is getting stuck exploring subregions of the map.
Finally, we inspected the trajectories of each agent when solving the task after 1 million training steps.
Figures 9 shows the DDQN agent and Figure 10 shows the DQRM agent. Both agents solved the
task, but DQRM solved it faster (83 steps vs 102 steps). In solving this problem, the main difﬁculty
for the DDQN agent is in reacting to entering the South room and discovering that the keys are not
there. Its reaction is to enter and leave the empty room many times before going to the North room.
In the case of DQRM, the agent goes directly to the North room after observing that the South room
is empty, but then it enters and leaves the North room a few times before collecting the keys. After
collecting the ﬁrst key, both agents solved the rest of the task almost optimally.
18
Kὑ1ὑ1
(a) Collecting the ﬁrst key.Kὑ1
(b) Opening the ﬁrst door.
Kὑ1
(c) Collecting the other key.K
(d) Opening the second door.
Figure 9: The learned trace after 1 million training steps by DDQN given a perfect RM, divided into
four stages.
Kὑ1ὑ1
(a) Collecting the ﬁrst key.Kὑ1
(b) Opening the ﬁrst door.
Kὑ1
(c) Collecting the other key.K
(d) Opening the second door.
Figure 10: The learned trace after 1 million training steps by QRM given a perfect RM, divided into
four stages.
19

Rodrigo Toro Icarte
Richard Valenzano
Ethan Waldie
Margarita P Castro
Toryn Q Klassen
Sheila A McIlraith

Typically in reinforcement learning, the agent only gets access to a reward signal: it sees a single number saying how well it has done. The problem might be simpler to solve if the agent could get a more holistic view of the problem through a structured representation of the reward. This could allow it to infer things like “if I went left, I would get 5 reward, but if I went right, I would get 10 reward”. Under the current RL paradigm, it has to try both actions in separate episodes to learn this.

Model-based RL tries to recover some of this structured representation: it learns a model of the world and the reward function, such that you can ask queries of the form “if I took this sequence of actions, what reward would I get?” The hope is that the learned models will generalize to new sequences that we haven’t previously seen, allowing the agent to learn from fewer environment interactions (i.e. higher sample efficiency).

This work does something similar using _reward machines_. The key idea is to represent both the reward and some aspects of the dynamics using a finite state machine, which can then be reasoned about without collecting more experience. In particular, given a POMDP, they propose learning a set of states U such that when combining the observation o with the state u, we have an MDP instead of a POMDP. This is called a _perfect_ reward machine. To make this feasible, they assume the existence of a labeling function L that, given a transition <o, a, o’>, extracts all of the relevant state information. (Since POMDPs can be reduced to belief-space MDPs, it is always possible to extract a perfect reward machine by having U be the set of possible beliefs and L be the identity function, but the hope is that U and L can be much simpler in most cases.)

They provide a formulation of an optimization problem over finite state machines such that a perfect reward machine would be an optimal solution to that problem (though I believe other imperfect reward machines could also be optimal). Since they are searching over a discrete space, they need to use a discrete optimization algorithm, and end up using Tabu search.

Once they have learned a reward machine from experience and a labeling function L, how can they use it to improve policy learning? They propose a very simple idea: when we get experience <o, a, o’>, treat it as a separate experience for every possible u, so that you effectively multiply the size of your dataset. They can then learn optimal policies that are conditioned on the state u (which can be inferred at test time using the learned state machine). Experiments show that this works in some simple gridworlds.

To summarize my summary, this paper assumes we have a POMDP with a labeling function L that extracts important state information from transitions. Given this, they learn a (hopefully perfect) reward machine from experience, and then use the reward machine to learn a policy more efficiently.

I see two main limitations to this method. First, they require a good labeling function L, which doesn’t seem easy to specify (at least if you want a high-level labeling function that only extracts the relevant information). Second, I think their heuristic of using every <o, a, o’> transition as a separate experience for every possible u would not usually work -- even if you learn a perfect reward model (such that the combination of o and u together form a “state” in an MDP), it’s not necessarily true that for every possible state in which you get observation o, when taking action a, you get observation o’. The authors acknowledge this limitation with an example of a gridworld with a button that changes how transitions work. But it seems to me that in practice, the underlying state in a POMDP will often affect the next observation you get. For example, in Minecraft, maybe you get some experience where you chop down a tree, in which your next observation involves you having wood. If you generalize it to all possible states with identical initial observations, you’d also generalize it to the case where there is an enemy behind you who is about to attack. Then, your policy would learn to chop down trees, even when it _knows_ that there is an enemy behind it.

It seems pretty important in RL to figure out how to infer underlying states when working in a POMDP, as it seems like a useful inductive bias for our agents to assume that there is a (Markovian) world “out there”, and I’m excited that people are thinking about this. Due to the two limitations above, I don’t expect that reward machines are the way to go (at least as developed so far), but it’s exciting to see new ideas in this area. (I’m currently most excited about learning a latent state space model, as done in e.g. <@Dreamer@>(@Dream to Control: Learning Behaviors by Latent Imagination@).)