Learning Reward Machines from Partially Observed Policies
Authors: Mohamad Louai Shehab
Antoine Aspeel
Necmiye Ozay
Venue: Accepted by TMLR ()

Abstract: Inverse reinforcement learning is the problem of inferring a reward function from an optimal policy {or demonstrations by an expert}. In this work, it is assumed that the reward is expressed as a reward machine whose transitions depend on atomic propositions associated with the state of a Markov Decision Process (MDP). Our goal is to identify the true reward machine using finite information. To this end, we first introduce the notion of a prefix tree policy which associates a distribution of actions to each state of the MDP and each attainable finite sequence of atomic propositions. Then, we characterize an equivalence class of reward machines that can be identified given the prefix tree policy. Finally, we propose a SAT-based algorithm that uses information extracted from the prefix tree policy to solve for a reward machine. It is proved that if the prefix tree policy is known up to a sufficient (but finite) depth, our algorithm recovers the exact reward machine up to the equivalence class. This sufficient depth is derived as a function of the number of MDP states and (an upper bound on) the number of states of the reward machine.{These results are further extended to the case where we only have access to demonstrations from an optimal policy. Several examples, including discrete grid and block worlds, a continuous state-space robotic arm, and real data from experiments with mice, are used to demonstrate the effectiveness and generality of the approach.

Published in Transactions on Machine Learning Research (10/2025)
Learning Reward Machines from Partially Observed Policies
Mohamad Louai Shchab
mishchab&Qumicch.com
Deportment of Robotics
University of Michigan, Ann Arbor, USA
Antoine Aspeel
Universite Paris-Saclay, CNRS, Centrale-Supéler Laboratoire des Signaur et Systèmes, Gif-sur- Siffsur-Yvette, France
Necmiye Ozay
necmige@xuiich.ete edu
Deportment of Electrionl Engineering and Computer Science Department Robotics
University of Michigan, Ann Arbor, USA
Reviewed on OpenReview: https: // openreview. net/ forum? id= 76bYYNoA7E
Abstract
Inverse reinforcement learning is the problem of inferring reward function from an optimal policy or demonstrations by an expert. In this work, it is assumed that the reward is expressed as reward machine whose transitions depend on alomic propositions associated with the state of Markov Decision Process (MDP). Our goal is identify the true reward machine using finite information. To this end, we first introduce the notion of prefix tree policy which associates distribution of actions each state of the MDP and each attainable finite sequence of atomic propositions. Then, we characterize equivalence class of reward machines that can be identified given the prefix tree policy. Finally, we propose a SAT-based algorithm that uses information extracted from the prefix tree policy to solve for a reward machine. It proved that if the prefix tree policy is known up to a sufficient (but finite) depth, our algorithm recovers the exact reward machine up to the equivalence class. This sufficient depth is derived as function of the number of MDP states and (an upper bound on) the number of states of the reward machine. These results are further extended to the case where we only have access to demonstrations from an optimal policy. Several examples, including discrete grid and block worlds, continuous state-space robotic arm, and real data from experiments with mice, are used to demonstrate the effectiveness and generality of the approach.
1 Introduction
Several frameworks exist for solving complex multi-staged tasks, including hicrarchical reinforcement learning (HRT.) (Pateria et al. 2021) reward machines (RMs) (Tcarte et. al. 2018) and linear temporal logic (L.TL.) al., 2021). HRL leverages a decomposition of tasks into subtasks, enabling agents focus on solving smaller, manageable problems before integrating solutions into higher-level policy (Sutton et al. 1999). On the other hand, RM and its generalizations (Corazza et all finite-state machines, capturing temporal dependencies and logical constraints in & concise and interpretable manner, similar to LTL. This structure simplifies policy learning and improves efficiency, especially in environments with long horizons or sparse rewards.
specifications (Chou al. 2020: Vaczipoor
2022) encode task-specific knowledge
As an extension inverse reinforcement learning (IRL) Ng and Russell 2000). one could ask the ques- tion of learning RMs, which enables agents autonomously extract structured representations of tasks, significantly enhancing their ability to solve complex, temporally extended problems. By learning reward

---

Published in Transactions on Machine Learning Research (10/2025)
machines directly from demonstrations, agents can adapt to tasks without requiring manually specified task representations, making this approach scalable and practical for real-world applications, such robotic ma- nipulation and autonomous vehicle navigation (Camacho et al. 2019 Icarte et al. 2023; Xu et al. 2020 Baert ct al. 2024 Camacho et al. h.21 This capability is especially valuable in environments where high-level lask-relevant features (aka, atomic propositions) are observable, underscoring the imporlance of learning RMs in advancing autonomous docision-making systems. For instance, in a high-level indoor navigation or patrolling task, semantic room labels can act as such propositions. Beyond autonomy applications, IRL and RM learning can also be used to infer the agent's s (e.g., humans") intentions to design incentives or better decision making environments (Nitschke et al. 2024 As will be demonstrated, such intent inference can also be used in neuroscience to analyze animal behavior and decision making.
Some of previous work on learning reward machines from data either assumes that the machine's states are observed (Araki et al.) 2019) or the reward is observed (Xu et al. 2020 Icarte et al.) 2023 Hu et al. 2024:
Abate et a 2023). . In the latter case the problem becomes finding a reward machine consistent with the observed inpul-oulput traces. Other work (Hasanbeig al.
2024
2021;
Furelos-Blanco
al.] 2020) infers
reward machines by combining automata synthesis with reinforcement learning and querying the environment for experiences. Others (Xu ct al.) 2021 Memarian et al. 2020) use the standard L algorithm for automata learning (Anglun 1987) to learn a consistent reward machine This assumes access to an oracle that can answer membership and conjectures queries. There are also works that only use observations of atomic propositions (Camacho et al. 2021; Baert et al. 2024), similar to us; however, they are limited to single- stage goal reaching tasks, where the RM has a simple structure that is merely used to obtain dense rewards. In parallel, several works aim to infer an L.T. specification from demonstrations satisfying and/or violating the specification (Neider and Gavran 2018 Chanuae-Chanlatte and Seshia 2020), requiring potentially large, labeled data set. Since LTL learning problem is inherently ill-posed, several regularization techniques are used such as formula templates or concept classes.
To the best of the authors' knowledge, no prior work has formalized and solved the problem of learning reward machines from partially observed optimal policies directly without the need to observe the rewards or the machine's state. The two main challenges of this setting are 1) partial observability (the reward is not observed, only the atomic propositions are observed). 2) partial reachability (not all transitions of the reward machine are visited in a given environment). In this work, address these challenges by first characterizing what can be learned in this setting (i.e., an equivalence class of reward machines) and then proposing a SAT-based algorithm, which provably learns a reward machine equivalent to the underlying true one. The key insight of our algorithm is to identify pairs of atomic proposition prefixes, namely negative ezamples, that lead to different nodes of the underlying reward machine from the observable optimal prefix-based policy, and encoding these examples as constraints in the SAT problem. We show that our method can be applied even when the optimal policy is accessible only through a finite set of optimal trajectories. To this end, we approximate the policy from the data and replace the SAT problem with variant called weighted MAX-SAT that provides robustness to incorrectly labeled negative examples. We demonstrate the efficacy of our algorithm in diverse settings, including grid-based MDPs, a robotic control task, and a real-world dataset of mouse navigation.
Given a set X, we denote by A(X) and |X] the set of all valid probability distributions on X and the cardinality of X, respectively. 1(X) denotes the indicator function of X. X*, denote the set of all finite/iffmnn /nfinite sequences of elements in X. For a sequence 7 and non-negative integers i, Ts denotes the 7t element of T; |r] denotes the length of Tend denotes the last element of 7 when is finite; Taj denotes the subsequence starting with the 8th clement and ending with the 5th; and denotes the subscquence ending with the tt element.
Notation:
2 Preliminaries and Problem Statement
2.1 Markov Decision Processes and Reward Machines
A Markov Decision Process (MDP) is a tuple M (S,A, Po. y,r), where S is a finite set of states, A is finite set of actions, P:S S x A A(S) is the Markovian transition kernel, ni A(8) is the initial state
2

---

Published in Transactions on Machine Learning Research (10/2025)
The set of distribution, ? E [0,1) is the discount factor and r 8 x A x 8 - R is the reward function.
feasible state trajectories for an MDP M, denoted by (MM,, is defined as:
(MM) = {{(0, $1, s1,...) ES" 33 e1....) E A P(sc+1 St. ae) > 0, Vt).
When we want to refer to finite prefixes of (MM, we simply use T.M and we omit M when it is clear from the context.
An MDP without the reward is referred to as an MDP model, and is denoted by M/r. MDP models can be decorated with labels. We denote such labeled MDP models as Mr = (S. P, No: 7, AP), where : - AP a labeling function that assigns to each state an alomic proposition, representing high- level conditions satisfied that state, from the set AP. A labeled MDP has corresponding language ((ML) (AP)", with ((M2)= = for E (AP)* o= I(),, where T E (T(M))): where we overload L to take in sequences. We also define the prefixes of language as:
(C) = {w E (AP)* Ez E L, s.t. w is a prefix of æ}.
The set of reachable states for a proposition sequence is:
(each(o) = {s E T T. s.t. L(r) = a, 7M = s}.
A Reward Machine (RM) is a tuple R = AP, 6,) which consists of a finite set of states U, an initial state ur U. an imput alphabet AP, a (deterministic) transition function du U x AP - U, and an output function & U x AP R. To avoid ambiguity between MDP states and RM states, the latter will be referred to nodes. The reward machine without the reward is denoted as G - R/6,, and we refer to it as reward machine model. We extend the definition of the transition function to define 6; : U x (AP)* + U as (1 (u, input words which can be used to reach w:
lo. ... , (x) = óu( -- (ōu (duu(,, 4), ... .4). Given a state u E U1 we define the paths of as the
(u) = (w E (AP)" | (, w) = =J.
We overload the operator Reach to include the set of MDP states reachable at u. It is given by:
Reach (u) = {s Reach(o) a Paths(u)).
As running example, we borrow the patrol task from [carre et al. 2018). Consider the room grid world shown in Figure Ta It is 4 by 4 grid where the agent move in the four cardinal directions, with a small probability of slipping to neighboring cells. We color-code different cells to denote the proposition label associated with the correspouding cell. For example, all cells colored green have the high level proposition A. The agent is tasked to patrol the rooms in the order A + B + - D. This is captured by the reward machine shown in Figure [b] Assume the agent's state trajectory starts with
+
(a,,a2, ag. b Ci). The proposition sequence associated with T a AAABC. The RM nodes traversed by following are (uo, u1, u1, M1, W2, us). Since e₁ is the only state than can be reached with a, we have that Reach(o) = {c1). Similarly, a E Paths(ug). Each transition in the reward machine gives reward of zero, except the transition from Uy to up i.c., (u, () = 1 if (u, I) (ux, D), and zero otherwise.
(2,
2.2 Markov Decision Process with Reward Machine
A Markov decision process with reward machine (MDP-RM) is a tuple RM (M/r, R, L) where M and
R are defined in Section 2.1 and L is labeling function L: S + AP. An MDP-RM can be equivalently seen as product MDP MProd (S', A', Pt, Mo. where S' = S x U, A' = A. P'(s', u' w'|s,, (s, u.aa) = P(s'|s, a)l(u'= 1 (u, L(s')), (0))),// 7/ =, 16 E ASS x U)) with u)( = Mo(s) =u) and r' a, a, s', , (u, M(8)). To make the notation compact, we denote the product state by == (8,u). The product of an MDP model with a RM model is a product MDP model G x Mr = (8', A', T', pb 7') defined similarly.
A trajectory of the product MDP MProd is a sequence (5g, ag. Su. au a1, ) where 5y = (0, uy) and ag 0. An initial state 8o is sampled from (ip- The introduction of 5g and ag at the start of the trajectory
3

---

Published in Transactions on Machine Learning Research (10/2025)
(-A,0) 0)
d,
d3
C1
C3
(4,0)
D
H
C
start
81
(-B,0)
d1
C2
(D,1)
(B,0)
H
H
H
a4
83
b
by
(-1,0)
ug
(C, 0)
HHE
B
a1
22
b1
b
(-C,0)
(a)
(b)
(c)
Figure 1: (a) The room grid world. (b) The patrol reward machine. (c) The room grid world with a hallway.
is to ensure that so induces transition in the reward machine. The reward machine thus transitions to = Su(s, I(so))- The agent then takes action and transitions to s1. Similarly, the reward machine transitions to リi = 6n,, L(s1)). The same procedure continues infinitely. We consider the product policy TProd: Domprod - ((A) where Domprod C x U is the set of accessible (s, u) pairs the product MDP. This policy a function that describes an agent's behavior by specifying an action distribution at each state. We consider the Maximum Entropy Reinforcement Learning (MaxEntRL) objoctive given by:
MMxxEm((TTTT') 1EST
(3,,, S:+1) AH((T(1 A
t=
where A > 0 is regularization parameter, and H(T(-(s) = = (ols)) is the entropy of the aed
policy л. The expectation is with respect to the probability distribution P the induced distribution over infinite trajectories following т, μo: and the Markovian transition kernel P' Ziebarr et al. 2008). The optimal policy "Prod; corresponding to a reward function r', the maximizer of i.c.,
Prod arg max JMaxEEt (x:r²).
(2)
Optimal product MDP trajectories are trajectories of the product MDP generated using FFrod- We overload this definition to optimal trajectories of the MDP, which generated from the optimal product MDP tra- jectories by simply removing the u states. For the rest the paper, optimal trajectories or demonstrutions refer the optimal trajectories of the MDP.
2.3 Prefix Tree Policy
Since the RM is unknown and the state of the RM is unobserved, we need representation of the agent's policy that is independent of the RM state u. We accomplish this by defining a prefix tree policy (PTP) the function that associates distribution over the actions to each state and each finite sequence of A(A), with
atomic propositions that can be generated by the MDP. It is denoted as ТРTР DompTP Domrrr ((s. a) E Pref([(MM.) and Reach(a)). An important remark here is that the agent is acting according to a policy TProd (als, u), since the agent has access to u. The PTP in turn encodes the information of the agent's product policy in terms of the variables that we have access to, namely the MDP states only. The relation between the two policies is governed by:
TPTP(a|s, o) (alod (ū
3
where a E (r) In particular, we say that the product policy xProd induces TPPP- We define the depth-l restriction of PTP as its restriction to the set (U) S x (AP)) n Domprr. That is the policy associated to words of length up to l. Tt is denoted by क|-p-
4

---

Published in Transactions on Machine Learning Research (10/2025)
The induced PTP captures both what is observable about product policy and what is reachable the product. Therefore, we can only learn reward machine up to whatever information is available in its induced PTP. We formalize this with the following definition.
Definition 1.
Two reward machines are poliey--quivvllet with respect a labeled MDP model if the product policies obtained by solving problem ( for each of the reward machines induce the some prefia tree policy defined in (3). Among all the reward machines that are policy equivalent with respect to labeled MDP, we define ménimal reward machine as one with the fewest number of modes.
Several equivalence relations among reward machines in the literature are special cases of policy equivalence. For instance, when learning finite state machines from observed rewards (Xu et al.] 2020 Giantamidis et al. 2021), two reward machines are said to be input-output equivalent if they produce the same reward sequence for the same atomic proposition sequence. Such input-oulput equivalent reward machines are clearly policy-equivalent. Furthermore, for trivial, i.e. one state, reward machine, our definition reduces to the policy-oquivvlence definition in the standard inverse reinforcement lcarning problem (Shchab et al.
2024
Cao et al. 2021
2.4
Problem Statement
Consider a labeled MDP model Mr. and a prefix tree policy true induced by an optimal solution of prob- We are interested in the following two problems in this paper:
lem 2
(P1) Docs there always exist depth-I" such that, given the labeled MDP model ML, a bound Umax _ttuu.l" of
on the number of nodes of the underlying reward machine, and the depth-I" restriction *PTP the true prefix tree policy, it is possible to learn reward machine that is policy-oquivalent to the underlying one?
(P2) If in problem (P1) exists, find a minimal reward machine that is policy-equivalent to the underlying one.
In what follows, we first provide an algorithm that takes the labeled MDP model Mr, the depth-/ restriction ttruel NPTP of the prefix tree policy TPTP for some arbitrary I, and the bound Umax and computes reward machine learned, .treel TPTP Then, we
that induces a prefix tree policy learned with the same depth-/ restriction, i.e., FTTP prove the existence of sufficient depth- in (PI), for which this algorithm solves problem (P2). We provide upper bound on 1" in terms of the number of states in the MDP and the number of nodes in the RM.
3 Methodology
3.1 SAT Encoding
We encode the RM lcarning problem into a Boolean Satisfiability problem (SAT). SAT is the problom of determining whether there exists an assignment to variables a given Boolean formula that makes evaluate to truc. While SAT is known to be NP-complete (Cook) 1971). there are several powerful off-tho-shelf solvers capable of solving large practical instances (Biere et al., 2009). Specifically, we use SAT to encode a graph with I Umax nodes and associate Boolean variable with each edge in the graph. Each node has [AP] outgoing edges. We define the Boolean variables {biks
1 1 i n, 1 S & z [AP]] as:
bikj = S1 if 5 j. to Otherwise,
(4)
where we use the shorthand 4 j to denote that proposition k transitions node i to node j. i.c., da(u; k) uj- We can encode several properties of the RM into Boolean constraints. Without loss of generality, we set node 1 of the graph to be uy. To make the derivation easier, we define for each atomic proposition i an adjacency matrix B& with (Bx)ss = baj- The Boolcan constraints we add are due to determinism, full-speciffcation,, negative ezamples and non- non-stutt of the learned reward machine. We expand on each of them below.
5

---

Published in Transactions on Machine Learning Research (10/2025)
Determinism: Due to the RM being deterministic machine, each label can only transition to one node. The corresponding Boolean constraints are:
Vi, k. Vÿ's }
bikj
biky' 0.
(5)
Full Specification: This constraint, also known as being imput-enabled (Hungar et al. 2003), ensures that all labels generate valid transitions at all nodes. The corresponding Boolean constraints are:
Vi, Vk, 3j such that būuj = 1.
(6)
We can combine the conditions of Sections 3.1 and 3.1
into one condition on each BN enforcing that each
row has exactly one entry with value 1.
Negative Examples: Our Boolean constraint here depends on the following result.
Lemma 1. Let a. E (AP)* be two finite label sequences. If (als, # ( o'), then 6ū a) # (1r,,, o').
Proof. It follows from
Based on this result, given the depth-l restriction TPTP of PTP иPтP, construct the set of negative examples as:
& 110,0'' merp(als,o) & xprp(als, for some s, s,a).
(7)
Let ........k ki and o' ki ki be two propositional prefixes that lead to different policies in the same state, therefore (o,0'} E,- We encode the condition given by Lemma into Boolean constraints as:
(BI, BT 4--
BT, e1) BTe))
A BI km
B e1) [0
0]"
where A is the element-wise AND operator and [1,0, .0 indicates that the paths start from the initial node. Our algorithm adds the Boolcan constraint given in S for each element of E The significance of encoding negative examples is that it eliminates the learning of trivial reward machines. In particular, our method never learns a trivial one-state reward machine with all trafsitaasssions (Icarte et al. 2023) as long as there is least one negative example in our prefix tree policy.
Non-Stuttering: A reward machine said to be non-stuttering if when a proposition transitions into node, that same proposition can not transition out of the node, i..,, for all (u,a, U AP U: 6u(u, a) = u' = (n(u',a) a) v'. This is related multi-stage lasks where the particular duration spent a subtask (i.e., satisfying given atomic proposition) is not important (Baier and Katoen 2008). When the reward machine is priori known to be non-stuttering, this extra condition can be included in the SAT problem. The main significance of this condition is trace-compression [Icarte et al. 2023), by which we can reduce the number of negative examples in -- by only keeping the shortest negative examples among the equivalent ones. In this case, two negative examples are equivalent between two pairs, the corresponding label sequences differ only by the same proposition repeated consecutively more than one time. We encode non-stuttering into Boolean constraints follows:
Vi,jk bikj =1
byrs =
(9)
The utility of this constraint. is demonstrated empirically in Section 5.1
3.2 Algorithm
To learn minimal reward machine from the depth-l restriction of prefix tree policy, we proceed as follows.

---

Published in Transactions on Machine Learning Research (10/2025)
We start with one node and increase the number of nodes until the following SAT problem is feasible when instantiated for graph with nodes:
SAT, 5
for all (a, E
(SAT)
By construction, this is guarantoed to be satisfiable for some n I Umax, upon which a reward machine model Glearued can be constructed from the satisfying assignment's Bu's. Then, compute the product MDP model MM/r)lameed glearned x .M. The learned product policy is constructed as follows. For each length word E (APY and for all Reach(or, define
" Prod
learned (a|s, glearned, (u,,, "PPP (als,
(a)s
(10)
where the transition function gleamed, is the transition function of glearned
The last step is finding the numerical values of rewards that render the product policy mppoaed optimal for the product MDP model (M/r))earuue This is standard IRL problem without reward machines where the special structure of the rewards on the product can be represented as features. We solve this step using the method developed in (Shehab et al., 2024). Featurization (see, (Shehab el al. 2024 Section 4)) is used to enforce that the reward function of the product can be written as r((s, u) a, (s uu)) = 6,(u, L(s')), with
u' = Oul, L()'). This IRL method gives us the corresponding output function 8, of the reward machine.
Remark 1. Although we construin the oulput funclion of the reward machine be of the form & : U x AP
I (leading to what is commonly known as simple reward machines [learte al.) (018)),, this does not limit the generality of our framework. The same procedure remains applicable in the case of dense output function of the form 8, : U AP + [S x Ax S + RR. We focus on the simple output function formulation in this paper to simplify both the presentation and the interpretation of the resulting reward machines.
The overall procedure is summarized in Algorithm !
Algorithm 1: Learning a Minimal Reward Machine from depth-l Restriction of a Prefix Tree Policy
Input: Depth-l prefix tree policy wpTP labeled MDP ML
Output: Learned reward machine Rlearned
1-1
while SAT, is infeasible do n ++++
Construct & using 7
[BalleT + SAT, solution
glearned < Construct_ RM [MM_modl((ff (Bx)AA)
(M/r)d +- glearned x ML
foreach (AP)' do
foreach E Reach(() do
Define product policy: #Prod
learned (als,, du (aarnod,"((1,00) aPT (als, σ)
true.
11 8 + IRL tt_ fo__extract extrrct_rrewrd[(fr learned (Prod (M/r) hearued learned)
12 return Rlearued (glearued, 8,)
is important to emphasize that up to line 7 in Algorithm 7 neither the oplimalily of the prefa tree policy nor the transition kernel P of the MDP is used. That is, our method learns a reward machine model glearned in model-free fashion as long as the prefix-tree policy induced by reward machine. This is further illustrated through an example in Appendiz C The optimality with respect the MazEnt objective in 7 and the transilion kernel only comes into play to extract the numerical reward values in lines 8-11.
Remark 2.
7

---

Published in Transactions on Machine Learning Research (10/2025)
3.3 Proof of Correctness
Let Glearned be the reward machine model extracted from the SAT solution, with gleamed being the associated transition function. The first property of our SAT solution is that it is consistent with any fixed depth-of the prefix tree policy. We formalize this in the result below.
Proposition 1. Given the labeled MDP model Mr. the depth-l restriction true. of the true prefer policy trrue, "PTP, and en upper bound umax on the number of nodes of the underlying reward machine, let glearned be the output of our SAT problem, and define learned to be the (infinite depth) prefix tree policy induced by glearned Then, the learned and the true prefiz tree policies have the same depth-l restriction, i.e., TPP
learned, true.l
TPTP
Proof. See Appendix A.1
While Proposition represents a desirable property of our algorithm, being consistent with the depth-1 restriction of the true prefix tree policy is in general not sufficient to be consistent with the true (infinite- depth) prefix tree policy (as required in Problem (P1) from Section 2.4 This is potentially problematic if the agent demonstrates unseen changes in its policy for prefixes longer than L. At the same time, it is not possible to run our algorithm with the unrestricted prefix tree policy *PTP because it would lead to an infinite number of negative examples, i.e, |&& = co. Consequently, the algorithm would not verminate. Fortunately, we can show that is large enough, then increasing will not change the satisfying assignments of the SAT problem.
Proposition 2. Given Mr, an upper bound Umax on the namber of nodes of the underbying reward machine, and the depth-l restriction upTP of some prefiz tree policy TPTP, where |S|ummax Then, {Ba)kel satisfying assignment for
IAPI is
SAT Uman
5
6 19 for all o' E 8
if and only if it is satisfying assignment for all j 2 1 for
SAT SAT
1 for all fa, E E 8
Proof. See Appendix A.2
Now, we present the main result of this section. Our result guarantees that given a sufficiently deep restriction of the true prefix tree policy, our recovered reward machine will be consistent with true infinite depth prefix tree policy. That is, our algorithm guaranteed to find a reward machine that policy-equivalent to the true reward machine.
Theorem 1. Given a labeled MDP model Mr and the depth-l restriction trued of a prefix tree policy induced by reward machine Rimee with at most umax nodes, if 1 2 |S u then the reward machine Rlearned returned by Algorithm 1 is policy-equuine equivalent to Rimee with respect to ML.
Proof. Follows immediately from Propositions and 2 Tn particular, we know that glearned is a solution of umm for all j 2 I, due to Proposition 2 Combined with Proposition this means that TPTP true.) for all
learned
#PTP
Remark 3. Note that, in practice, a depth-l restriction where < S| can be sufficient to find a reward machine that is policy equivalent to the true one if all the solutions of the corresponding SAT problem are equivaeett to each other (e.g., they correspond to the same reward machine up to renaming of nodes). This will be further illustrated in the experiments of Section 5.
4 Learning From Demonstrations
In this scction, we present how our method can be applicd when the optimal policy is known only through set of optimal demonstrations D ((0,, 81, )2) Hence, the depth-l restriction of the true prefix tree
8

---

Published in Transactions on Machine Learning Research (10/2025)
policy is unknown. Consequently, we construct the following consistent unbiased estimate:
Σ 1= l1 = 0,8; -, -.
(11)
E I(o 0,
However, the fact that we only have access an approximation of TerP leads two challenges. First, using RPTP to directly construct the set of negative examples & will lead to many pairs {a,a'} being incorrectly considered negative examples (see equation ( Second, the SAT problem is not robust to incorrect negative examples. In the next subsections, we address these two problems.
4.1 Estimating Negative Examples
To limit the inclusion of incorrect negative examples in & pair of prefixes (01, will be considered as negative example only if we have high confidence that they correspond to negative examples given the true prefix tree policy. In particular, let 01, 72 be two prefixes from D that we want to compare and let E S be a state. For j E {1,2}, let ny be the number of visitations to the pair (5,oy), ie., the denominator in [11]. Note that Tppp(als, σ,) is a categorical distribution of which we have a sample estimate for. Hence, it follows from (Weissman
2003) that for all 0:
P P(Iftrrr(ass,o) Trrr ( az)ll s0) 85,
(12)
=nee" We pick probability of the two confidence intervals described in [12 (for (1,2)) to overlap is lower bounded by
with 6, A (21-4 2n,e,
|TPP
# 02)]1I. Using a union bound, the
1 1-81 6g. Consequently, 81
consider the pair {o1, 02) being negative example if there is state such where a is user-defined parameter.
that 1 - 81 - 82 2
4.2 MAX-SAT: SAT with Robustness to Incorrect Negative Examples
Even when the set of negative examples is constructed from pairs of prefixes for which we have high confidence of being true negative examples, some pairs may still be mislabeled as negative examples. We refer to those as false positives. To deal with these false positives, we implement a weighted MAX-SAT (Manquinho et al. 2009; Biere et al. 2009) variant of the SAT problem. Concretely, the weighted MAX- SAT problem consists of finding a Boolean assignment of the variables diky such that (i) constraints (5). 6 and hold, i.e., the recovered RM is deterministic, non-stuttering, and fully specified; and (ii) the constraint
8 holds for a maximum number of pairs {o,o') in & In other words, (5), 6 and (9) are hard constraints whereas the constraints coming from S are soft constraints with equal unitary weights. This approach allows for robustness to misclassifications of negative examples. We empirically demonstrate in Section 5.3 that this weighted MAX-SAT formulation consistently infers the smallest set of consistent reward machine models, even in the presence of false positives.
5 Experiments
To demonstrate the generality and efficiency of our approach, we apply to diverse set of domains, from classical grid-based MDPs to continuous robotic control task and real-world biological navigation dataset. In every experiment, we fix the discount factor to = 0.99 and the regularization weight to X = 1.0 when solving Problem [ both for generating demonstration traces and for reward recovery. Our code is imple- mented in Python, and the Z3 library [De Moura and Bjørner 2008)) is used for solving the SAT and weighted MAX-SAT problems. To enumerate all the satisfying assignments of the SAT problem, we add a constraint that the next solution should be different every time our SAT solver finds solution. Our implementation code is made publicly available here: https:/ https://gituub.om. hwtp://mllhhhhhmmmnenaa seetaaninnrewwarrdaaeer machincs.hee......
5.1 Tabular GridWorld MDPs
Our first experiment (patrol) on the muning-exampll patrol task of Figure | bound from Theorem 1 is 256. However, with only depth-6 prefix tree policy, end up with total of
By setting umax 4, our
9

---

Published in Transactions on Machine Learning Research (10/2025)
6 solutions. These are all the possible renamings of the true reward machine (see Figure 1b meaning that the true reward machine is learned up-to-renaming with a depth-6 prefix tree policy. Table summarizes the findings, where we also show how the non-stuttering condition of Section 3.1 helps reduce the size of the negative example set, yet still rocovering the true reward machine model. While some reduction in the SAT solver time achieved, the drastic gain is the time required to encode all the negative examples inlo the SAT solver, making the overall procodure orders of magnitude faster. For the remaining experiments, non-stuttering is assumed.
Task patrol with patrol without 19)
Depth
[S]
|
31
SAT time (sec)
(9)
16
4
3076
0.51
16
4
30573
1.73
patrol-hallway
9
4
241435
2.859
Table 1: Solution statistics for the tabular GridWorld MDP.
For our second experiment (patrol-hallway), we add a hallway between the rooms, as shown in Figure Ic This is a 3 x grid world, where the corresponding label of each room is shown. The reward machine is kept the same, and the atomic proposition H is added as self-loop to all the nodes, With this added hallway, longer atomic proposition prefixes are required to reach all nodes of the reward machine, showing how the underlying MDP affects the required depth for learning a reward machine. For example, the shortest atornic proposition that can reach U3 is now AHBHC instead of ABC as in the previous example. With a depth-9 prefix tree policy, the reward machine is again learned up-to-renaming. Additional experiments studying transfcrability and the sensitivity to the upper bound Umax can be found in appendices
D and E
5.2 Tabular BlockWorld MDPs
The setup for this experiment is a modified block world problem (Khodeir al., 2023 Wolfe and Barto experim 2006). There three blocks colored green, yellow and red, as well as piles. Each stacking configuration of the blocks is a state of the MDP, and the action space consists of selecting pile to pick from and a pile to place onto. We can only grab the top block of any stack. Action outcomes are assumedd deterministic. The goal in the first task (stack) is to stack the blocks in the ordered stacking configurations st1, st2, stg, shown in Figure 2a All other states have the label i, denoting intermediate states. The corresponding reward machine is shown in Figure 2b the robot stacks the blocks in the order st, + st2 + sta, it gets a reward of 1.0. With a depth-10 10 prefix tree policy and Umax 3, our algorithm recovers consistent reward machines, which are the true reward machine up-to-renaming. The findings summarized in Table Results obtained from finite set of demonstrations are reported in Appendix B.1.1
st 1
st
stz
(-st, 0)
(-st2, 03
sts O}
(sta
(sttg.
(staa: 1.0)
(b)
Figure 2: (a): Block World MDP. The left-most stacking configuration has label st1, where all blocks are stacked on the first pile with green being under yellow and yellow being under red. Similarly, the middle configuration has label sta and the right-most configuration has label st&. (b): Stacking Reward Machine.
For our second task (stack-avoid), we introduce "bad" state, shown in Figure 3a The true reward machine is shown in Figure 3b] The robot's task is to stack the blocks in the order st, - st2 without going through stpa. through stpa, it will get a smaller (yet more immodiate) reward of 0.2, but will get stuck at u3 with 0 reward forever. We note that the product policy is uniformly random in both Uy and Ug. This means that proposition traces such st,, i, sta and st,, i, stbd sta look identical from policy perspective, both
does so, it reaches U2 and gets a reward of forever. If during execution it passes
10

---

Published in Transactions on Machine Learning Research (10/2025)
reach nodes with uniformly random policies, while the first being more desirable than the second. By setting Umax = 3, depth-8 policy was sufficient to recover the reward machine shown in Figure 3e up-to-renaming: that is, we find a smaller reward machine consistent for this task. Notably, setting Umax = 4 with the same depth policy yields more than a 1000 solutions, indicating that uniquely recovering the same sized reward machine -as the ground truth machine- requires longer depth policies. Results obtained from a finite set. of demonstrations are reported in Appendix B.1.2
(-sti. Dj
(-st2,
(True, 1)
(st 0)
()
sty 6.44) ( (-st2 -alaa, 5.79)
(True, 6.45)
(st,, 12.35)
(st 8.63
stbd
(slad 0.2)
(etba, 0.2)
start
(stta 0)
(True,
(sbba, 6.57)
(b)
(c)
Figure 3: (a) The block stacking configuration with label stad that we want our robot to avoid. (b) The ground truth reward machine. (c) Smaller consistent reward machine with the task stack-avoid.
Task
Depth
S
A
SAT time (sec)
stack 10
60
9
73548
0.612
stack-avoil
8
60
9
24763
0.569
Table 2: Solution statistics for the tabular BlockWorld MDP.
5.3 2-Link Robotic Arm with Continuous State Space
Our setup for this experiment is modified Reacher-v5 environment (Gymnasium (Towers et al. 2024 MuJoCo [oooooo et al. 2012) ) where planar robotic must reach targets randomly placed within the arena. The agent's 's state is 10-D continuous vector consisting of the end-effector's position and velocity, the target's coordinates, and their mutual distance. Actions are continuous torques in []1 applied at the elbow and shoulder joints.
Inspired by (Arki et al. 2021), we fix three colored targets blue (b), red (r), and yellow (y) in distinct quadrants of the arena (see Figure 4). The goal is reach them in the order b - y -> r. To train an expert policy, we discretize each torque dimension into five levels {-1.0,-0.. 0.0, 0.5, 1.0) (yielding 25 total actions). We then employ Proximal Policy Optimization (PPO) (Schulman all 2017). ; implemented in Stable-Baselinee (Ralfin et. al.] 2021 to maximize reward given by the negative Euclidean distance between the end-effector and the active target.
When generating trajectories, we switch the desired target immediately upon reaching the previous target, thereby emulating reward machine without. introducing the machine during training or simulation. Ap- pendix B.2 provides full discrctization and training dotails. From this procodure we collect 1M trajectorics.
Table 3 reports, for various probability thresholds (defined in Section 4.2 the size of the negative-example set E and its False Positive Rate ( As expected from Section 4 solving SAT problem with all the negative examples yields 0 solutions. Instead, we employ our weighted MAX-SAT variant and recover two viable solutions (Figures 4b and 4cb up-to-renaming, which represent the best achievable given our dataset. These are the same solutions that we recover if instead we supervised with the ground tarnh reward machine to remove the false positives, emphasizing the robustness of this approach.
IPPR simply the number of false positives divided by |-
11

---

Published in Transactions on Machine Learning Research (10/2025)
b
y
T
-
b
y
b
start
u0
U1
U2
start
u0
U2
(b)
(c)
Figure Reacher experiment. (a): 2-link robotic arm with the three colored targets. (b) First recovered
reward machine model. (c) Second recovered reward machine model.
D
[
7
3
FPR
# Weighted MAX-SAT solutions
1M
17.4M
25
160
1472
0.001
1.90%
1M
17.4M
25
160
1193
0.0001
1.67%
4
1M
17.4M
25
160
882
0.00001
1.36%
4
Table Summary of the experiments for the 2D-contimoous reacher arm.
5.4 Real-world Mice Navigation
We also applied our lcarning framework to the trajectorics of real mice navigating (in the dark) in a 127-node labyrinth maze [Rosenberg et al. 2021) shown Figure 5 Each node, labeled with number Figure 5a represents a state the MDP. The mouse can select from actions: (stay, right, left, reverse). State 116 (middle-right) contains water resource and labeled w, state 0 (center) is the home state labeled h, and all other states are labeled i. Two cohorts of 10 mice moved freely in the maze for 7 hours, with cohort being water-restricted and the other was not. A sample water-restricted mouse trajectory is shown in Figure 5b) Difference in water restriction condition resulted in different animal behavior between these two cohorts. For the purposes of our study, we only considered the water-restricted mice. We used the same dataset of trajectories from (Ashwood et al.) 2022), which is comprised of 200 mouse trajectories, given as state-action pairs of length each. We set our probability threshold a = 0.001. With Umax = 2, our algorithm learns the unique reward machine shown in Figure 5c This is consistent with the seen behaviors of water-restricted mice which first. try to reach the water port and hydrate, after which their behaviors switch to exploring the
maze going back to the home state (Zhu et al. 2024 Ashwood et al. 2022). More details can be found in Appendix B.3.
Exit
(h, 0.00)
(h, 0.62)
(w, 12.10
start
up
U1
(w, 0.12)
(i. 0.877
(i, (1.0.86)
Entry
(b)
(c)
Figure Labyrinth experiment. (a): Maze structure and state space definition. (b): Trajectory of single mouse. (c): Recovered reward machine. (a) and (b) are reprinted from (Rosenberg et al. 2021
We further evaluate the quality of the recovered reward machine using a held-out set unseen trajectories (20 test trajectorics). The reward machine model and the product policy are Iearned from the remaining 180 training trajectories. For each trajectory the test set, we compute its lokelioel under the product policy. This metric reflects how well the learned model captures the underlying trajectory distribution
12

---

Published in Transactions on Machine Learning Research (10/2025)
and generalizes beyond the training data. The results are shown in Table and represent the average log-likelihood over the test dataset. To contextualize the results, we compare against three baselines: a uniformly random policy and two variants of Max Causal Entropy IRL Ziebart (2010). D-IRL employs a dense feature representation, assigning a distinct reward value to every MDP state. F-IRL uses a structured Teature representation aligned with the environment's 's labeling. Specifically, its feature vector is one-hot encoding indicating whether the agent is at the home port (h), at the water port (w), or any other intermediate state (i). These results also quantitatively demonstrate the superiority of the learned reward machine in capturing the unseen behaviors of the mice.
Method
D-IRI.
F-IRI.
Uniform
LRM
Average log-likelihood
-28.99
28.57
-30.9
9.81
Table 4: Average log-likelihood performance of the learned reward machine on unseen trajectories, compared against three baseline methods.
6 Limitations and future work
The present framework assumes that the atomic propositions provided by the labeling function exact, meaning no noise or mislabeling is present in the observed labels. While this assumption simplifies the analysis and ensures that the negative examples used in the SAT formulation are correct, it may be unrealistic in settings where the labeling function is derived from perception modules or noisy sensors. Future work could address this limitation by extending the framework to handle noise the labeling function, for instance by integrating robust RM-learning techniques such as those in Parac ct al. (2024) thereby broadening applicability real-world domains.
Another limitation arises from the requirement to identify all negative examples up to depth / = |8] [S] for the theoretical results to apply (sce Theorem D. Although this bound is polynomial in the size of the MDP and the maximum number of RM nodes, the number of distinct negative examples -and thus the number of clauses in the SAT formulation may grow exponentially with / in the worst case. For large problems, this can lcad to SAT instances of prohibitive size and render the approach computationally intractable. On the other hand, in practice, we were able to identify the reward machine model with a much smaller depth in all our examples. Therefore, the existence of tighter sufficient depth bound is open problem, which will investigate further in the future. We are also interested in developing smart prefix-selection algorithms that incorporate a verification loop to check, on the fly, whether the selected negative examples are sufficient.
7 Conclusion
In this work, we present framework for learning reward machines from partially observed policies, where neither the rewards nor the reward machine states are available to the learner. Instead, our method constructs a SAT problem instance from a sufficiently deep prefix tree policy, from which a reward machine equivalent to the true underlying one can be recovered. The approach is generalized to learn the reward machines directly from demonstrations from an optimal expert, where robustness to inacouracies policy estimates is handled using finite-sample confidence bounds and MAX-SAT used to eliminate incorrect negative examples. A diverse set of experiments illustrate the effectiveness of our approach.
Acknowledgments: This work is supported in part by ONR CLEVR-AI MURI (#N00014- 21-1-2431).
13

---

Published in Transactions on Machine Learning Research (10/2025)
References
Alessandro Abate, Yousif Almulla, James Fox, David Hyland, and Michael Wooldridge. Learning task automata for reinforcement learning using Hidden Markov Models. In European Conference on Artificial Intelligence, pages 3-10. IOS Press, 2023.
Dana Angluin. Learning regular sets from queries and counterexamples. Information and computation, 75 (2):87-106, 1987.
Brandon Araki, Kiran Vodrahalli, Thomas Leech, Cristian-loan Vasile, Mark D Donahue, and Daniela L Rus. Learning to plan with logical automata. Robotics: Science and Systems Foundation, 2019.
Brandon Araki, Xiao Li, Kiran Vodrahalli, Jonathan DeCastro, Micah Fry, and Daniela Rus. The logical options framework. In International Conference on Machine Learning, pages 307-317. PMLR, 2021.
Zoe Ashwood, Aditi Jha, and Jonathan W Pillow. Dynamic inverse reinforcement learning for characterizing animal bchavior. Advances in Neural Information Processing Systems, 55:29663-29674 2022.
Mattijs Bacrt, Sam Leroux, and Picter Simoens. Reward machine inference for robotic manipulation. arXiv
preprinl ar Xrv:24 12. 10096, 2024.
Christel Baier and Joost-Pieter Katoen. Principles of model checking. MIT press, 2008.
Armin Biere, Marijn JH Heule, Hans van Maaren, and Toby Walsh. Handbook of satisfiability, volume 185. TOS press, 2009.
Alberto Camacho, Jacob Varley, Andy Zeng, Deepali Jain, Atil Tscen, and Dmitry Kalashníkov. Reward machines for vision-based robotic manipulation. In IEEE International Conference Robotics and Au- tomation, pages 14284-14290. IEEE, 2021.
Alejandro Camacho, Carlos Toro learte, Tor Klassen, Richard Valenzano, and Sheila A Mcllraith. LTL and beyond: Formal languages for reward function specification in reinforcement learning. In International Joint Conference on Artificial Intelligence, pages 6065-6073, 2019.
Haoyang Cao, Samuel Cohen, and Lukasz Szpruch. Identifiability in inverse reinforcement learning. Advances in Neural Information Processing Systems, 34:12362 12373, 2021.
Glen Chou, Necmiye Ozay, and Dmitry Berenson. Explaining multi-stage tasks by learning temporal logic formulas from suboptimal demonstrations. In Robotics: Science and Systems, 2020.
Stephen A Cook. The complexity of theorem-proving procedures. Proceedings of the third annual ACM symposium Theory of computing, pages 151 158, 1971.
Jan Corazza, Ivan Gavran, and Daniel Neider. Reinforcement learning with stochastic reward machines. In
AAAT Conference on Artificial Intelligence, volume 36, pages 6429 6436, 2022.
Leonardo De Moura and Nikolaj Bjørner. 73: An efficient SMT solver. In International conference on Tools and Algorithms for the Construction and Analysis of Systems, pages 337- 340. Springer, 2008.
Daniel Murelos-Blancoo Mark Law, Alessandra Russo, Krysia Broda, and Anders Jonsson. Induction of subgoal automata for reinforcement learning. In AAAI Conference on Artificial Intelligence, volume 34, pages -890- 2020.
Georgios
Giantamidis, Stavros Tripakis, and Stylianos Basagiannis. Learning Moore machines from input- output traces. International Journal on Software Tools for Technology Transfer, 23(1): 2021.
Hosein Hasanbeig, Natasha Yogananda Jeppu, Alessandro Abato, Tom Melham, and Daniel Kroening. Sym- bolic task inference in deep reinforcement learning. Journal of Artificial Intelligence Research, 80:1099- 1137, 2024.
14

---

Published in Transactions on Machine Learning Research (10/2025)
Mohammadhosein Hasanbeig, Natasha Yogananda Jeppu, Alessandro Abate, Tom Melham, and Daniel Kroening. DeepSynth: Automata synthesis for automatic task segmentation in deep reinforcement learn- ing. In AAAI Conference on Artificial Intelligence, volume 35, pages 7647- 7656, 2021.
Jueming Hu, Yash Paliwal, Hyohun Kim, Yanze Wang, and Zhe Xu. Reinforcement learning with predefined and inferred reward machines in stochastic games. Neurocompating, 599:128170, 2024.
Hardi Hungar, Oliver Niese, and Bernhard Steffen. Domain-specific optimization in automata learning. In International Conference on Computer Aided Verification, pages 315-327. Springer, 2003.
Rodrigo Toro Icarte, Toryn Klassen, Richard Valenzano, and Sheila McIlraith. Using reward machines for high-level lask specification and decompasition in reinforcement learning. In International Conference on Machine Learning, pages 2107-2116, PMLR, 2018.
Rodrigo Toro Tcarle, Toryn Q Klassen, Richard Valenzano, Margarita P Castro, Ethan Waldie, and Sheila A McIlraith. Learning reward machines: A study in partially observable reinforcement learning. Artificial Intelligence, 323:103989, 2023.
Mohamed Khodeir, Ben Agro, and Florian Shkurti. Learning to search in task and motion planning with streams. IEEE Robotics and Automation Letters, 8(4):1983-1990, 2023.
Vasco Manquinho, Joao Marques-Silva, and Jordi Planes. Algorithms for weighted boolean optimization. In International Conference on Theory and Applications of Satisfiability Testing, pages 495-500 508. Springer, 2009.
Farzan Memarian, Zhe Xu, Bo Wu, Min Wen, and Ufuk Topeu. Active task-inference-guiieed deep inverse reinforcement learning. In IEEE Conference on Decision and Control, pages 1932 1938. IEEE, 2020.
Daniel Neider and Tvan Gavran. Learning linear temporal properties. In 2018 Formal Methods in Computer Aided Design, pages 1-10. IEEE, 2018.
Andrew Y Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning, volume 1, page 2, 2000.
Paul Nitschke, Lars Lien Ankile, Eura Nofshin, Siddharth Swaroop, Finale Doshi- Velez, and Weiwei Pan.
AMBER: An entropy maximizing environment design algorithm for inverse reinforcement learning. In ICML Workshop on Models of Human Feedback for Al Alignment, 2024.
Roko Parac, Lorenzo Nodari, Leo Ardon, Daniel Murelos-Blanee Federico Cerutti, and Alessandra Russo. Learning robust reward machines from noisy labels. arXiv preprint arXiv:2408. 14871, 2024.
Shubham Pateria, Budhitama Subagdja, Ah-hwee Tan, and Chai Quek. Hierarchical reinforcement learning: A comprehensive survey. ACM Computing Surveys, 54(5):1135, 2021.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of machine learning research, 22(268):1-8, 2021.
Matthew Rosenberg, Tony Zhang, Pietro Perona, and Markus Meister. Mice in a labyrinth show rapid learning, sudden insight, and efficient exploration. Elife, 10:c66175, 2021.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimiza- tion algorithms. arXiv preprint 170: 06347, 2017.
Mohamad Louai Shehab, Antoine Aspeel, Nikos Aréchiga, Andrew Best, and Necmiye Ozay. Learning true objectives: Linear algebraic characterizations of identifiability in inverse reinforcement learning, In 6th Annual Learning for Dynamics & Control Conference, pages 1266-1277. PMLR, 2024.
Richard S Sutton, Doina Precup, and Satinder P Singh. Between MDPs and Semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181122 211, 1999.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IEFE/RSJ International Conference on Intelligent Robols and Systems, pages 5026-5033. TEEE, 2012.
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulao, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv22477.10032 2024.
Pashootan Vaezipoor, Andrew C Li, Rodrigo A Toro Icarte, and Sheila Mcilraith. LTL2Action: Gen- eralizing LTL instructions for Multi-Tan Task RL. In International Conference on Machine Learning, pages 10497-10508. PMLR, 2021.
Marcell Vozquez-Cannlate and Sanjit. A Seshia. Maximum causal entropy specification inference from demoustrations. In 32nd International Conference on Computer Aided Verification, pages 255-278. Springer, 2020.
Tsachy Weissman, Erik Ordentlich, Gadiel Seroussi, Sergio Verdu, and Marcelo J Weinberger. Inequalities for the 11 deviation of the empirical distribution. Hewlett-Packard Labs, Tech. Rep, page 125, 2003.
Alicia P Wolfe and Andrew G Barto. Decision tree methods for finding reusable mdp homomorphisms. In National Conference Artificial Intelligence, volume 21, page 530, 2006.
Zhe Xu, Ivan Gavran, Yousef Ahmaa,, Rupak Majumdar, Daniel Neider, Ufuk Topcu, and Bo Wu. Joint inference of reward machines and policies for reinforcement learning. In International Conference on Automated Planning and Scheduling, volume 30, pages 590-598, 2020.
Zhe Xu, Bo Wu, Aditya Ojha, Daniel Neider, and Ufuk Topcu. Active finite reward automaton inference and reinforcement learning using queries and counterexamples, International Conference on Machine Learning and Knowledge Extraction, pages 115-135. Springer, 2021.
Hao Zhu, Brice De La Crompe, Gabriel Kalweit, Artur Schneider, Maria Kalweit, Ilka Diester, and Joschka Boedecker. Multi-intention inverse q-learning for interpretable behavior representation. Transactions on Machine Learning Research, 2024.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum cansal entropy.
Carnegie Mellon University, 2010.
Brian D Ziebart, Andrew I, Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse rein- forcement learning. In AAAI Conference on Artificial Intelligence, volume 8, pages 1433-1432 Chicago, II, USA, 2008.
16

---

Published in Transactions on Machine Learning Research (10/2025)
Appendices
A Proofs
A.1 Proof of Proposition
[
Proof. We proceed by contradiction. Assume that there exists some o' E (ref(((M))) with |o S l, a
state E Reach(o') and an action A such that:
learned (als, + TPTP (als, " "PTP
(13)
Let 8ū learned (ur, o'). The hand-ide of 133 can be written as:
"PTP (als
(als, .), for some E Path(u) n (Ce())))))) with |a| S learned (als, u) = TPPP (
Prod
where the second equality is due to 100. We get:
PPT (als, σ # TPTP (als, where both o, E Path(u).
More precisely, du learued (u,, .) = oü learned 0'). Due to the contrapositive of Lemma we have a con- tradiction. Similarly, due to the full specification condition of Section 3.1 and (10). the support of the two prefix policies will be the same by construction.
A.2 Proof of Proposition
We start by formally defining some important concepts that will be central for proving our result. Our proof idea requires reasoning about joint (i.e. synchronized) paths over two distinct reward machine models, and being able shrink these joint paths by removing cycles (i.e. loops). To start, we define cycles in product MDP model G x Mr, follows:
Definition Given product MDP model G x Mr and proposition sequence == -...4 generated from slate sequence T (2, ................................ 88)) (i.e., Pref (re)()(M)))), say that subsequence ouj of is cycle in G Mr if s, = and (,) = (3),)
We will the above definition to construct shorter label sequences with no cycles given long label sequence. In particular, let le & { be the number of states in G x ML. By the pigeonhole principle, we know that any state trajectory of length more than le has to visit some product state more than once, meaning that it has at least one cycle. In particular, given any proposition sequence a, with |o| > les let ō the subsequence of obtained by removing all the cycles in a. Then, we know that |o] S le, since ā has at most all the unique states from a, which cannot exceed le- Note that removing cycles preserves the last product state reached from following a.
Next, define synchronized reward machine models.
Definition 3. Let 91 (41 u), AP, G2 (Ug, AP, 82) be two reward machine models, with Unl. [122 Umnxx- The synchronized reward machine model is reward machine model defined follows:
Geyue (Usvue ugyne .AP, 88yuc)
yyyuc Ui x U2,
uf sync (u},u?), geyac (w 12). (6, (,,,,)) IEAP.
Similarly to a regular reward machine model, the product gavne Mr well defined. The total number of states in gayne Mr is upper bounded by I [ In particular, consider an arbitrary label sequence E Pre)()(MM))). generated from state sequence (81, 82, ".,.""",8"""" ,8), with lol > [M]u
17

---

Published in Transactions on Machine Learning Research (10/2025)
and let 22) u²) ou (u)"m,,) ,0). This means that the synchronized product state (u, (u,"',s") is reachable in gayue x MM. Thus, by removing cycles of a in Geyue x MM. we can construct a shorter prefix o, with o S MM such that (4, ou buyne (uym (uy .), and s* is the MDP state reached.
We are now ready for the proof Proposition
2
We provide Figure 6 as an illustration of the proof.
Proof. For a shorthand notation, denote the first SAT instance the proposition statement as SAT and the second SAT instance as SATumee.j Also note that & E when j 2 I as the negative examples only grow the depth increases. Throughout the proof, we interchange TPTP and TPTP both are equal up to depth j.
We need to show that the additional negative examples in & \ & do not change the set of satisfying assignments. The = direction is straightforward, since a satisfying assignment cannot become unsatisfying by removing constraints.
|AP| for SAT and assume by contradiction that [Bal==1 is not satisfying for SATumax. with j > I. Consider the reward machine model glearned, with the transition function glearned, corresponding to {Bx JAPl Since J
For the = direction, take a satisfying assignment {Ba)k= |AP]
IAPI is unsatisfying for SATumaxdh then
there exists negative example fa, E \ ET such that
glearned, (uz, σ) = glearned, slearedd"(u,,""") (u,.o'),
(14)
while TaTp(a"|s", (a"|s", o) + "ppppa"'s" o') for some (s",a") A. These two facts are shown as Appre and =u connecting a and o' in Figure B
Now, let giree (U1 uy, AP, gtrue) be a reward machine model consistent with TPTP. We define the following nodes, along with the associated product states:
true strue (ur, (uerue, gicue ML
W learned
A glearned o), (u learned , E glearned x ML srue, =strue (u),o'), o'), (u "ruo, s") E gtrue MM. (u,, learnedy ,8) E glearned ML.
W learnod, glearned
Hence, by 14
we have that learned u learned,
Let gayne be the synchronized reward machine model
between G glearned and Grrue according to Definition
3 We observe the following:
(u learned) 8 (eyyuc,(1" (un .)) (utrue learned :s)) E gayne x ML.
(u tre. / learned. 5ayme, (u, sync oa), (utrued, /earned,), ,,* E Gayne Mr.-
This means that the synchronized product states (utrue, u learned s*) and (utrued, wlarodd ,s*) ) are both reachable in Gayne x M,. Thus, by removing cycles, there must exist shorter sequences, o,ō', with [o], |ō'l S |S w such that:
(utree u learned) =6 gayne," (1 sync :), (u truey u Jearned./) -5yyne,+ (unn, (ugync ,
(15)
Note that s" is still the reached MDP state in both synchronized product nodes above. By the definition of 01
caync,* we can decompose 15
into:
u true (in F), u learned = glearned (in ( σ),
true./ strue (urue(ur, (ur, σ'), d /earnedd/ 8 glearned (u),
(16)
This means that and lead to the same node truu in girue, Similarly, a' and o' both lead to the same node trruey Since gimee is consistent with TPTP, the following holds:
TpTp(a|s", o) ( .), Va E
TPTp(a|s",o) = Va e A.
(17)
18

---

Published in Transactions on Machine Learning Research (10/2025)
Since TPIP(a"|s" 00) # TPrp(a"|s* .o) due to contradiction assumption, conclude from [7
that TPTP(( (a" 1 5) * mPTp(a")> ,')) However, dearned wlearned, combined with (16) implies that
glearne (u,,ō) o) glarnod o¹), contradicting that (Ba]k=1 is SAT assignment for the depth I. This concludes the proof.
TPPP
TPPP
HPTP
TPPT
j
Figure 6: Proof Illustration of Proposition Kprps Apppp means that the prefix tree policy лPTP is equal/different for the corresponding sequences. means that the corresponding sequences arrive at the node in glearned
Experiment Details and Additional Results
B.1 Tabular BlockWorld MDP
To generate the expert trajectory dataset, we first compute the product policy given the BlockWorld MDP dynamics and the corresponding ground-truth reward machine. We use soft bellman iteration to find the optimal product policy. The procedure is sunmarized in Algorithm ?
At each time step, use the visited proposition prefix to find the reached node the true reward machine, using u ()) We then sample action from the true product policy x <a(als, u), where is the MDP state reached. We keep count of the sampled actions at all state-prefix pairs for constructing the prefix tree policy. The procedure summarized in Algorithm 3 In the following section, we show the results for the stack and stack-avoid tasks.
B.1.1 (stack) task
We simulated varying size demonstration datasets and solved our SAT problem with each. In Table 5 we show the effect of the number of demonstrations in reducing the number of satisfying solutions. With a demonstration datasct of size 1M, the ground-truth (stack) reward machine is recovered up-to-renaming. However, even a demonstration set of size 5000 can indeed reduce the number of satisfying solutions to These 4 solutions include the ground-truth reward machine (accounting for solutions due to renaming), and the reward machine (up-to-renaming again) shown in Figure 7b where 6n a) u1, instcad of 8, 8u(u2, a) U2 as in the ground-truth reward machine.
B.1.2 (stack-avoid) task
In Table B we similarly show the effect of the size of the simulated dataset on the number of satisfying solutions for the stack-avoid task. In this task, 1M trajectories were not enough extract all the necessary negative examples to reduce the number of solutions to 2. However, much smaller dataset of size 1000
19

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm 2: Soft. Bellman Iteration the Product MDP
Input: MDP M, ground-truth reward machine R, labeling function L, tolerance c.
Output: Optimal product policy n² Construct product MDP MMrod (S', A', P', plo y', r') as in Section 2.2 Initialize
Vo(o, - 0 v(s,u) E ES' error 00, k+ 0
while error > e do foreach (s, E S' do
5
foreach a E A' do
L Qs+1 (8, t (() + o Ecoranes ((s,, (8', ,)) Vi((d, .d)
6
V1+1(s, + log Σ exp(Qx+1(s,,,,0)
7
aed'
enror & (.) (s, Vi(ss u)
k t k 1
10 foreach u) E S', a A' do
x <a u) + exp exp(Qu(s,u,a) - vu(s,u))
11
12 return n
Algorithm 3: Constructing the Prelix-Tres Policy via Simulation
Input: MDP M, reward machine R, true product policy r" 1s u), # trajectories N, trajectory length
Output: Prefix-tree policy -H Initialize counts: C(a, s, +- 0 for all prefixes a, states E S, actions A. for i +-1 to N do
σ - (empty prefix)
s ~~ Ho (initial MDP state)
u + dy (initial RM node) for 1 H do u (,,,,u,
Sample ~ n² s,
Execute a, observe s' ~ P(s, a) and label E f I(8');
10
a+
xtolle olle (concatenation);
11
C(a,s,a) + 1(;, s';
12
13 foreach prefix o, state 8, action do
(a (8)0) = C(a,8,a)
14
C(a,sa)
return TPTP
20

---

Published in Transactions on Machine Learning Research (10/2025)
[
[3
# SAT solutions
1000
20
58
0.05
24
3000
20
174
0.05
12
3000
20
305
0.05
8
10000
20
490
0.05
8
0.1M
20
2249
0.05
4
1M
20
7995
0.05
2
Table 5: Size of the simulated dataset vs. the number of satisfying solutions.
st
stz
sts A
stt
sts
st2
stg
st1
stz
start
do
U1
U2
st
st2
start
ug
U1
U2
st1
st3
stg
(a)
(b)
Figure 7: Recovered Reward Machine Models for stack task.
enough to reduce the number of solutions to 4, which we manually inspect. These solutions are shown in Figures Sa and 8b (up-to-renamingg. The reward machine in Figure 8b differs from the ground- truth reward machine by assigning (u(u1,stbdl = u1, instead of 6u(u1, stba) u2. Such an ambiguity could have been resolved by a negative example of the form {aı = st,, st, 02 = st,, i,s st which the MDP is capable of producing given the length of the simulated trajectories. However, since the policy becomes uniformly random after reaching stad, simulating such negative examples might require much larger dataset.
D
3
# SAT solutions
200
20
26
0.05
32
500
20
52
0.05
1000
20
67
0.05
0.1M
20
1615
0.05
1M
20
5330
0.05
Table 6: Size of the simulated dataset the number of satisfying solutions.
st,,
sta A
stad
st1
sts
True
st
stz
st
st2
start
u1
start
up
U2
st stbd
stbd
stp
(a)
(b)
Figure Recovered Reward Machine Models for stack-avoid task.
21

---

Published in Transactions on Machine Learning Research (10/2025)
B.2 2-Link Reacher Robotic Arm with Continuous State Space
B.2.1 Expert policy training using deep RL
To train the policy, we used the standard Reacher-v5 environment from Gymnasium, augmented with cus- tom wrappers to discretize the action space and introduce randomized imitial conditions. We also increased the episode horizon from the default 50 steps to 160 steps in order to match the expected time to visit all 3 de- sired poses and finish the task. The siscreteReacherrctioiiii action space to a MultiDiscrett space of 25 torque combinations by restricting each component to one of five values: {-1.0, -0.5, 0.0, 0.5, 1.0). Additionally; the weweeaanndmzzzdRRaccciie randomized initial joint positions and velocities at each episode reset to promote robustness and generaliza- tion. This modification was essential, as the default initial conditions in the base environment are typically limited to small perturbations around fully-cxtended arm configurations. Training solcly from such narrow initializations led to poor task performance and limited the policy's 's ability to generalize across the broader state space.
maps the continuous two-dimensional
wrapper introduces
We trained a Proximal Policy Optimization (PPO) agent using the stable-baselinees library, with all the default settings and hyperparameters. To accelerate training, we employed vectorized environment with 50 parallel instances running on CPU. The agent was trained for 20 million time steps using multilayer perceptron (MLP) policy.
B.2.2 Discretizing the state space and sampling the policy
The MDP state space is generated by discretizing the shoulder angle 01, the elbow angle 022 and the angular velocities 01, .2. The angles are discretized uniformly in [-x, л] with bin size of 10 so 0.17 radians. The angular velocities are discretized uniformly in [-14, 14] with a bin size of 0.25 rad/s. This leads to z 17.4M states in the MDP.
When sampling an action from the optimal policy, we first identify the discrete state corresponding to the current continuous state. The action then sampled at the center point of that discrete region. This approach slandardizes the policy's behavior across different continuous states that map the same discrete representation, ensuring consistency within each discretized region of the state space.
B.3 Real-world Mice Navigation
For recovering the reward function, first need to construct the product policy. However, real-world data generally do not perfectly fit the assumed mathematical models, ie., the max entropy assumption. Since most of the states in the MDP are unvisited, and many action are not sampled at most of the states, the product policy is predominantly sparse (~~ 92% of the action probabilities are 0). We preprocess the policy by clipping its minimum value to 0.05. This allows us to use the results from (Shehab et al. 2024) to compute the reward function, as we require computing the log of the policy. The reward space constrained to satisfy the feature-dccomposition which the reward machine provides. However, there is generally no guarantee that we can find such leaturized reward function that can perfectly induce any product policy. With the shown reward in Figure 5c the norm difference between the induced policy and the true policy is 2.66 (the norm difference with a uniformly random policy is 2.97). By increasing Umaa to 3, one of the recovered reward machines is shown in Figure g which yields a slightly larger policy difference norm of 2.71.
C Policy Optimality is not Required
In this section, we illustrate with a toy example that the optimality of the prefix tree policy is not necessary to recover the reward machine model. We consider the MDP given by the BlockWorld of Section 5.2 and a simple 2-node reward machine model given by [0] We do not specify a reward function here nor solve the forward MaxEnt RL problem before simulating the dataset, and instcad specify an arbitrary stochastic policy at each node of the reward machine. With dataset size 50K and trajectory length 20, our method learns uniquely the reward machine model.
22

---

Published in Transactions on Machine Learning Research (10/2025)
st
st2
(h, 0.003
(w, 0.111
(h,0.a 0.333
12.701
(h, 0.84)
start
No
U1
U2
(1, (1.0.11)
st
start
Uo
U1
(1,0.87)
(i. 0.877
(w, 0.93)
st²
Figure 9: Recovered Reward machine for Mice
Dataset with Umax 3.
Figure 10: Reward machine model considered
in Appendix C
D Sensitivity to the Node Bound Umax
To understand the effect of misspecifying the number nodes on the resulting learned reward machine, we run an additional experiment the Gridworld benchmark (sco Section 5.1 while varying the bound on the number of nodes. We simulate dataset of 2, 500 trajectories of length 10 each. We then construct the negative example set using finite sample results and solve for reward machine model using MaxSAT. The recovered models are shown in Figure [1 below. These reward machine models still qualitatively capture the overall task even with the limited number of nodes. To quantitatively assess the quality of these models, below. We then perform trajectory rollouts to
we first learn & product policy for each using Algorithm g compute the return of the policy consistent with each model. While the actions are sampled according to the corresponding lcarned policy of cach model, the reward calculated based on the traversod nodes the ground-truth reward machine, given by Figure [PP Results in Table show graceful degradation in performance the number of reward machine nodes decreases. We also implemented Max Entropy IRT. baseline Ziebart al. 20008). where static reward is learned from the demonstration set.
B
D
D
B
A
B
start
40
u1
start
d0
U1
V/2
D
(a) Mmnax 2
(b) Mmax = 3
Figure 11: Recovered reward machine models with varying node bounds.
Model
Total # clauses
# used clauses
Rollouts
|
Average Return per Rollout
Umux 4
1112
1112
10K
100
176.62
Umax 3
1112
902
10K
100
148.74
Umax
=2
1112
589
10K
100
94.25
D-IRL
10K
100
17.12
F-IRL
10K
100
14.49
Table 7: Statistics for the varying bound experiment. (Total # clauses) the number of negative examples from the finite dataset. (# used clauses) is the number negative examples which MaxSAT included as hard constraints. (Rollouts) is the total number of rollouts. ([7l) is the length of the trajectory. (Average Return per Rollout) is the sum of accumulated rewards in all the rollouts divided by the number of rollouts.
For this experiment the reward for completing the task 10.
23

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm 4: Construct Learned Product Policy
Input: Learned reward machine model G, Ground-truth reward machine model G True prefix tree policy SPTP, Labeling function L, MDP M, prefix length dp
Output: Learned product policy ă
forcach node u E 9 do
Compute Reach(u);
// we consider prefixes up to a fixed length of dp
foreach Paths(u) do
Compute Reach(o);
foreach E Reach(o) do foreach E A do
7
i(a s, + (arr(a s,
Normalize 7 and fill in zero rows with uniform policy;
return π
E Transferability Performance
In this section, test the transferability of our recovered reward machines when deployed different euvironment with the same labeling function. In particular, the new environment is similar to the Tabular GridWorld MDP studied in Section 5.1 except that the room assignment is changed, as shown in Figure 12a below. The dynamics are kept the same. order to test transferability, we first generate trajectory dataset using the ground-truth optimal policy in the new environment. Then, for each reward machine model, given by a varied Umax; we construct. a product policy according to Algorithm | (using the prefix tree policy of the original environment) and solve for reward function that will be optimized in the new environment. We then compute the loellikel of the trajectory datasct given the corresponding optimized policies. The results are shown in Figure 12b] where (TRUE) is the log-likelihood of the ground-truth optimal policy in the new environment. It can be seen that the performance of the learned reward machine with 4 nodes almost matches that of the ground-truth reward machine. The results also display a graceful degradation in performance with decreasing number of nodes.
b4
by
a4
ag
IRI-F
-90.9
IRL-D
-19.7
b1
b2
a
a2
UNIFORM
20.75
15.76
C4
C3
d,
d3
TRUE
C₁
C2
di
d2
24
22
20
18
-1
-12
10
(a)
(b)
Figure 12: (a): Modified labeled MDP. (b): Average logelikel of different models over 100 trajectories generated by the optimal policy in the new labeled MDP.
24


Reviews:
[{"review_id": "ZYQWVIUukO", "reviewer": "Reviewer_Q4TZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n1. The paper is very well written, and the reward machine setting is clearly defined.\n2.  Rigorous proofs are given for Theorem 1 which justify the equivalence between the learned policy and the policy of the reward machine\n3. This paper strictly improves on previous work, so it is a generally important/helpful line of research. \n\nWeaknesses: \n\n1. The setting cannot handle (and does not address) vector rewards\n2. The setting only considers deterministic reward machines", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "At a high level, this paper identifies a sufficient condition toward inverse RL: the problem of learning the reward function given the optimal policy (where we only observe the atomic proposition associated with the reward function), in the setting where the reward function is given by a reward machine whose transitions depend on certain prefix trees.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "-", "broader_impact_concerns": "-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0ZVknTGtQ3", "reviewer": "Reviewer_Rfpt", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is the first to address the problem of learning reward machines when only a partially observed policy is available, filling a gap in prior work that assumed access to rewards or full policy information.\n- It provides formal guarantees that, under specific assumptions, the method will recover a reward machine equivalent to the ground truth, giving the approach a solid theoretical foundation.\n- The work clearly separates the process of learning the machine’s structure from the process of estimating the numerical reward parameters, which makes the method modular and easier to adapt.\n- In practice, the approach is made more robust by using a MAX-SAT formulation that can tolerate noisy or imperfect data in the demonstration setting.\n\nWeaknesses\n- The method assumes that the set of atomic propositions and the labeling function mapping MDP states to these propositions are completely known and correct, which may be unrealistic in many applications.\n- It requires the reward machine to be deterministic and to have a defined transition for every possible proposition in every state, which are strong assumptions.\n- The approach depends on the observed policy being generated by a maximum-entropy optimal policy; deviations from this assumption could reduce identifiability and lead to incorrect machines.\n- The theoretical bound on the prefix length needed for guaranteed recovery is very large, making it impractical for large problems, and the paper does not provide a practical way to decide when enough data has been collected.\n- The noise handling in the demonstration setting relies on heuristic filtering and weighted MAX-SAT, without formal guarantees on accuracy in the presence of limited or noisy data.\n- The paper’s writing and presentation are not very clear, the lengthy definitions with dense symbols are hard to follow without explanation of the intuition or implication of them\n- Most experiments are verification-style, where the authors design environments with a known reward machine, generate data from it, and check whether their method recovers it; this mainly tests correctness rather than usefulness in real-world or transfer scenarios.\n- The only real-world experiment, using mouse navigation data, is interpreted qualitatively rather than validated quantitatively, so it does not fully demonstrate the method’s effectiveness when the ground truth is unknown.\n- The paper does not discuss related work in imitation learning or inverse reinforcement learning in sufficient depth, nor does it compare the proposed approach against these baselines. While the focus here is on recovering reward machine structure, it would be valuable to situate the method more clearly within the broader IL and IRL literature and to highlight its advantages over existing approaches.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies how to learn reward machines when you cannot directly observe the rewards or the machine’s internal states, but only see part of the optimal policy. The authors introduce the idea of a prefix tree policy, which records what the agent would do in a given MDP state after seeing a certain sequence of atomic propositions. They define what it means for two reward machines to be equivalent in terms of the policies they produce, and they show how to recover a minimal equivalent machine. The main technical tool is a SAT formulation that searches for a machine structure consistent with the observed policy. They prove that if you have a long enough prefix tree policy, the method can recover the exact equivalent machine. They also extend the approach to work from demonstration data by filtering likely errors and using a MAX-SAT variant. Experiments on grid worlds, block worlds, 2D reacher, and a real mouse navigation dataset show the method can find the correct machine up to renaming.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The paper is difficult to follow, especially for audience without much background of reward machines. The authors should make the writing clearer, motivate the formulation of the problem first, and add intuitive motivations/explanations/examples alongside the formal descriptions. \n2. The method assumes that the atomic propositions and labeling are completely correct. The authors should discuss how robust the approach is to incorrect or noisy labels and, if possible, test it on data with imperfect labeling.\n3. The experiments mostly test recovery in known settings. The authors should also show results in cases where the true reward machine is unknown or where the learned machine is applied to a different environment with the same propositions.\n4. The mouse navigation experiment is only analyzed qualitatively. The authors should add some form of quantitative evaluation for this case, even if it is indirect, to strengthen the real-world evidence.\n5. Relevant baselines, including standard IRL methods and Behavioral Cloning, should be discussed in more depth and compared against in the experiments to facilitate understanding the impact of this work.\n\n1-4 is important for my recommendation, 5 is nice to have.", "broader_impact_concerns": "No ethical concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "yXzg1WpBMA", "reviewer": "Reviewer_zSDg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strenghts**:\n\n   * **Novel and Important Problem Formulation**: The paper tackles a well-motivated problem. Learning structured\n     rewards from policy observations alone is a significant step beyond traditional IRL methods that rely on reward\n     observations or full state-space trajectories. This has interesting implications for autonomous systems and behavior analysis.\n   * **Theoretical Rigor**: The work is theoretically sound. The formalization of the Prefix Tree Policy and the characterization\n     of the RM equivalence class provide a solid foundation for the problem. The main theoretical result (Thm 1), which\n     guarantees the recovery of an equivalent RM, is a powerful and relevant statement.\n   * **Nice Algorithmic Approach**: The use of a SAT-based formulation is a nice tweak. It turns a complex\n     search problem into a well-understood constraint satisfaction problem, for which highly optimized solvers exist. Encoding\n     policy differences as \"negative example\" constraints is a particularly insightful way to prune the search space.\n   * **Robustness and Practicality**: The extension to learning from finite demonstrations via weighted MAX-SAT is crucial for\n     real-world applicability. \n   * **Satisfactory Empirical Validation**: By testing the framework on discrete, continuous,\n     and real-world biological data, the authors effectively demonstrate the versatility and robustness of their approach. The\n     successful application to the continuous Reacher arm and the real-world mouse data is particularly nice to me.\n\n**Weaknesses**:\n\n   * **Scalability Concerns**: The primary weakness, common to many SAT-based approaches, is scalability. The theoretical bound for\n     the required policy depth ($l > |S| u^2_{max}$) can be very large for complex environments. Furthermore, the number of\n     potential negative examples can grow exponentially with this depth, potentially making the SAT instance intractable for\n     problems with large state spaces or highly complex RMs.\n   * **Dependence on $u_{max}$**: The algorithm requires an upper bound on the number of RM nodes as an input. While\n     Algorithm 1 iterates the number of nodes n upwards, the theoretical guarantees are framed in terms of a known u_max. In a\n     real-world setting where the true RM is unknown, choosing an appropriate u_max (or knowing when to stop the iteration of\n     n) is a non-trivial challenge. An insufficient u_max could lead to the recovery of a non-equivalent, overly simplistic RM.\n\n   * **Assumption of Deterministic RMs**: The framework assumes that the underlying task is representable by a deterministic RM.\n     While this is a powerful and common model, some complex, stochastic, or history-dependent tasks might not be perfectly\n     captured by this structure.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper addresses a problem in the field of Inverse Reinforcement Learning (IRL): learning a structured reward function, represented as a Reward Machine (RM), from an expert's policy without access to direct reward signals or the internal states of the RM. The core contribution is a framework that can infer the structure of the underlying RM by only observing the expert's actions and the resulting sequences of high-level \"atomic propositions\" from the environment.\n\nThe  methodology involves three key steps:\n   - **Prefix Tree Policy** (PTP): A representation that maps a history of atomic proposition sequences (a prefix) and a current state to a distribution over actions. This captures how an expert's policy changes based on task progress.\n   -  **Equivalence and SAT Formulation**: They formalise the problem by defining a policy-equivalence class for RMs, as multiple RMs can produce the same observable policy. The search for a consistent RM within this class is then encoded as a Boolean Satisfiability (SAT) problem. Constraints for the SAT solver are derived from the structure of RMs and, crucially, from \"negative examples\", namely pairs of proposition sequences that lead to different observed policies, implying they must terminate at different nodes in the true RM.\n   - **Learning from Demonstrations**: The framework is extended to the more practical setting of learning from a finite set of expert demonstrations. To handle the noise and uncertainty inherent in policy estimates from finite data, the SAT problem is relaxed into a weighted MAX-SAT problem, providing robustness against incorrectly identified negative examples.\n\n\n  The authors provide a theoretical proof that their algorithm can recover a policy-equivalent RM given a policy known up to\n  a sufficient, finite depth. They validate their approach on a diverse set of experiments, including discrete grid worlds, a continuous robotic arm simulation, and a real-world dataset of mouse navigation, demonstrating its effectiveness and generality.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Low-Level**\n- (major) I would strongly suggest to provide references and background for SAT and (minor) to frame sub subsections as paragraphs, for readability.\n- Figure 1(c) is kind of far from its actual reference, I would suggest to adjust for this.\n\n**High-Level**:\n- I would strongly suggest to add in the introduction a high level description with examples of the setting of interest, naming some examples of what solving for such task might be, for instance by describing the experimental problems at hand.\n- It would be beneficial to add a more explicit discussion on the scalability limitations of the\n      approach. \n- Some guidance on setting the key. hyperparameter α (the confidence threshold for negative examples) would be valuable. A brief discussion on the sensitivity  of the results to this parameter would improve the paper's practical utility.\n- The paper mentions using weighted MAX-SAT to handle false positives but I did not find details on how the weights are assigned. A brief sentence explaining that  more explicitly would be nice. \n- A dedicated section on \"Limitations and Future Work\" would be a valuable addition.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "T6hsVi4oME", "reviewer": "Reviewer_EAZD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### **Strengths**\n\n* **Solid Theory:** It proves an identifiability bound: a prefix-tree policy (PTP) observed to depth\n  \\$l^{\\star}=|S|,u\\_{\\max}^{2}\\$\n  suffices to reconstruct the ground-truth RM up to policy-equivalence.\n\n* **Robustness to Demonstration Noise:** Treating negative-example clauses as *soft* weights in a MAX-SAT variant allows correct recovery even when some demonstrations are mislabeled.\n\n* **Diverse & Scalable Experiments:** Across four domains—grid-patrol, block-stacking, a continuous 2-D Reacher arm (1 M trajectories), and 200 real mouse-maze runs—the method attains exact or near-exact recovery with shallow depths and sub-second solver times.\n\n---\n\n### **Weaknesses / Points for Improvement**\n\n* **Strong Optimality Assumption:** All guarantees rely on having the exact optimal policy; real-world demonstrations are often sub-optimal or noisy, weakening theoretical assurances.\n\n* **Potentially Large Depth Bound:** The worst-case requirement \\$l^{\\star}=|S|u\\_{\\max}^{2}\\$ scales quadratically with RM size and linearly with the MDP’s state count—prohibitive for complex tasks even if practice needs less.\n\n* **Need for an Upper Bound \\$u\\_{\\max}\\$:** Both the theory and Algorithm 1 assume a tight prior on the number of RM nodes; guidance on selecting or adapting this bound is missing.\n\n* **Scalability of the SAT Core:** The negative-example set can explode into hundreds of thousands of clauses, leading to many SAT calls; further heuristics or incremental solving strategies are only briefly mentioned.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper addresses the open problem of recovering an unknown reward machine (RM) purely from observations of an optimal policy in a labeled MDP, without ever seeing numeric rewards or RM states.\nIt introduces the prefix-tree policy (PTP) abstraction, which records the action distribution for every state * proposition-prefix pair observed from the expert; this is the only information available to the learner in their “partially observed” setting.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### **Requested Changes**\n\n* **Include quantitative baselines.**\n  Add head-to-head comparisons with representative automata-learning (e.g., DFA inference) and inverse-RL algorithms (Bayesian IRL, AIRL, etc.) on at least one shared benchmark, reporting accuracy, sample efficiency, and runtime. Direct evidence of practical advantage is essential to justify acceptance.\n\n* **Relax or test the exact-optimality assumption (Critical).**\n  Either extend the theory to $\\varepsilon$-optimal / noisy experts or demonstrate empirically how performance degrades with sub-optimal demonstrations. Without this, the guarantees rest on an unrealistic premise for real-world data.\n\n* **Release anonymised code and data at review time (Critical).**\n  Providing the promised Python + Z3 implementation, solver settings, and datasets now would let reviewers (and readers) verify claims and reproducibility.\n\n* **Report sensitivity to the upper-node bound $u_{\\max}$ (Strengthening).**\n  Show how synthesis quality and runtime vary when $u_{\\max}$ is misspecified, and propose an adaptive or heuristic selection rule (e.g., MDL, BIC, incremental growth).\n\n* **Analyse the worst-case depth bound $l^{\\star}=|S|u_{\\max}^{2}$ in practice (Strengthening).**\n  Offer tighter bounds for common RM structures or empirical depth-versus-error curves to clarify scalability on larger tasks.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)