Active Reward Machine Inference From Raw State Trajectories

Abstract: Reward machines are automaton-like structures that capture the memory required to accomplish a multi-stage task. When combined with reinforcement learning or optimal control methods, they can be used to synthesize robot policies to achieve such tasks. However, specifying a reward machine by hand, including a labeling function capturing high-level features that the decisions are based on, can be a daunting task. This paper deals with the problem of learning reward machines directly from raw state and policy information. As opposed to existing works, we assume no access to observations of rewards, labels, or machine nodes, and show what trajectory data is sufficient for learning the reward machine in this information-scarce regime. We then extend the result to an active learning setting where we incrementally query trajectory extensions to improve data (and indirectly computational) efficiency. Results are demonstrated with several grid world examples.

• The paper introduces a SAT-based method that jointly infers reward-machine transitions and state labels from raw state trajectories without observing rewards, propositions, or automaton nodes.
• Theoretical analysis shows that a trajectory depth of |S|u_max² is sufficient for identifying a policy-equivalent model, while experiments recover benchmark machines exactly up to state and label renaming.
• An active query strategy reduced stored trajectories from 382,000 to 10,300, negative-example memory from 24.76 GB to 0.147 GB, and mean runtime from 7,185.5 to 3,544.8 seconds.

Problem setting and motivation

Reward machines (RMs) encode the memory structure of multi-stage tasks as finite-state automata whose transitions are triggered by high-level propositions, and they enable policy synthesis via product MDPs. Existing RM inference methods, however, universally assume access to some symbolic layer: observed rewards, ground-truth labels, machine nodes, or at least a predefined (possibly noisy) labeling function mapping low-level states to atomic propositions. This paper by Shehab, Aspeel, and Ozay removes that assumption entirely: the only inputs are an MDP model and a depth-limited history policy — a function from state trajectories to action distributions induced by the optimal product policy. Neither rewards, nor labels, nor automaton nodes are observed, and both the labeling function $L$ and the transition function $\delta_{\mathbf{u}}$ must be learned jointly.

The problem is inherently ill-posed because multiple labeled RMs can induce the same history policy; the paper therefore targets policy equivalence up to renaming, with minimality defined as fewest nodes. Two questions are posed: (P1) whether a finite sufficient trajectory depth exists given a node-count upper bound, and (P2) how to recover a minimal policy-equivalent labeled RM.

SAT-based joint inference of transitions and labels

The core technical device is a Boolean encoding of $\delta_{\mathbf{u}}$ and $L$ via binary variables $b_{jpi}$ (RM transitions per proposition) and $\mathbf{L}_{pk}$ (state-to-proposition assignment), with constraints enforcing well-definedness, anchoring ($u_I = 1$, $L(1)=1$), compatibility with negative examples, and an optional non-stuttering constraint for duration-insensitive multi-stage tasks.

Negative examples are derived from a simple but powerful lemma: if two state trajectories induce different actions at some state under the history policy, then their label sequences must lead to different RM nodes. Each such pair yields the SAT constraint $v_\tau \neq v_{\tau'}$, where $v_\tau$ is computed by Boolean matrix products over the encoded transitions and labels. Notably, only negative examples constrain the problem — positive evidence is implicit in satisfiability.

Sufficient depth result. The main theoretical contribution is Proposition 1: with node upper bound $\delta_{\mathbf{u}}$0, the depth

$\delta_{\mathbf{u}}$1

is sufficient — the SAT problem is satisfiable at $\delta_{\mathbf{u}}$2 if and only if it is satisfiable for all greater depths. The proof proceeds by contradiction using a synchronized labeled RM model combining the true and candidate machines: any violating negative-example pair longer than $\delta_{\mathbf{u}}$3 can be cycle-reduced in the synchronized product MDP (which has at most $\delta_{\mathbf{u}}$4 states) without changing either machine's node, producing a shorter counterexample that contradicts the depth-$\delta_{\mathbf{u}}$5 solution. This gives a finite identifiability guarantee in a regime where no prior work applies, since no symbolic observations exist. The output function $\delta_{\mathbf{u}}$6 is not addressed here; its recovery reduces to a structured IRL problem handled in prior work once the RM model is known.

Active extension of the history policy

Encoding all negative examples up to depth $\delta_{\mathbf{u}}$7 is computationally prohibitive because the number of state trajectories grows exponentially with depth. The key observation is that many extensions are redundant: if a length-$\delta_{\mathbf{u}}$8 trajectory reaches the same RM node as its length-$\delta_{\mathbf{u}}$9 prefix under the ground truth, it introduces no new negative examples. Since $\delta_{\mathbf{u}}$0 and $\delta_{\mathbf{u}}$1 are unknown, this motivates a volume-removal active learning strategy.

The algorithm maintains a candidate set of feasible SAT solutions, subsamples $\delta_{\mathbf{u}}$2 hypotheses, generates candidate trajectory pairs via randomized DFS targeting pairs predicted to end at the same node, scores each pair by how closely it bisects the candidate set (maximizing $\delta_{\mathbf{u}}$3), queries the top $\delta_{\mathbf{u}}$4 pairs against the true history policy, adds discovered negative examples incrementally, and re-solves the SAT instance without re-enumerating the full burn-in set. Termination occurs when the solution set collapses to a single model up to renaming.

Experimental results

Experiments use stochastic grid worlds with two tasks: pick_n_drop (pickup → drop-off cycling while avoiding a danger zone) and patrolABCD (patrolling four rooms in order). Ground-truth labelings are hidden from the algorithm.

For pick_n_drop, a depth-9 history policy recovers the ground-truth RM exactly, with all 12 remaining solutions differing only by renaming. For patrolABCD, depth 9 yields 36 solutions (6 node permutations × 6 label permutations), again matching ground truth up to renaming; a Tetris-shaped variant gives identical results.

The active extension results are strong:

Metric	Exhaustive	Active extension
Max depth reached	9	13
Trajectories stored	382K	10.3K
Negative examples	~414M	0.292M
Memory for negative examples	24.76 GB	0.147 GB
Mean total runtime	7185.5 s	3544.8 s

The exhaustive baseline becomes practically infeasible beyond depth 10 due to exponential branching (~1.53M trajectories at depth 10, growing roughly $\delta_{\mathbf{u}}$5 per step), whereas the active method holds the branch count near 10K via the query budget $\delta_{\mathbf{u}}$6. On pick_n_drop with $\delta_{\mathbf{u}}$7, convergence to the ground-truth solution set occurred in 100% of trials by depth 12, while the random-query baseline found essentially no constraining negative examples even at depth 20. On patrolABCD, $\delta_{\mathbf{u}}$8 converged in 96.6% of trials by depth 13 (83.3% for $\delta_{\mathbf{u}}$9), versus a random baseline still averaging 378 candidates with nearly two orders of magnitude larger variance. The runtime gain is nearly $L$0 despite reaching a deeper horizon, driven mostly by reduced SAT burden (2417 s vs. 7159 s mean).

These numbers support the paper's claim that selective querying alleviates both memory and computational bottlenecks; the implication is that RM inference without any labeling oracle is tractable on tabular problems of moderate size, which exhaustive approaches cannot handle past shallow depths.

Limitations and open questions

The authors are explicit that the framework is a foundational identifiability and algorithmic result rather than an end-to-end practical system. Several assumptions bear directly on the results:

• Discrete tabular setting: the SAT encoding requires finite $L$1, $L$2, and AP; scaling to perceptual state spaces would require learning labelings from raw representations.
• Exact history policy access: the theory assumes the true induced history policy; in practice it must be estimated from finite state-action data, and robustness to estimation error is deferred to statistical extensions of prior work rather than analyzed here.
• Termination criterion conservatism: termination requires uniqueness up to renaming, which may be stricter than necessary — earlier stopping would be possible if all candidates were verified policy-equivalent, and integrating such an equivalence test is left open.
• Node and proposition bounds: correctness relies on an upper bound $L$3; the effect of choosing it too small or too large is not studied empirically.
• Output function: reward recovery is delegated to prior structured IRL work; joint treatment under label uncertainty is not addressed.

A further open question raised is the relationship between task-induced memory (RMs) and memory arising from partial observability (information states, POMDP learning), toward a unified account of memory requirements in sequential decision-making.

Conclusion

This paper establishes that reward machines and their labeling functions can be jointly identified from raw state trajectories alone, given a depth-$L$4 restriction of the induced history policy, via a SAT formulation built on negative behavioral examples. An active extension algorithm that queries maximally hypothesis-splitting trajectory pairs achieves exact recovery up to renaming on grid world tasks while reducing negative-example memory by two orders of magnitude and total runtime by roughly half relative to exhaustive search. The guarantees hold for discrete MDPs with exact policy information; extending them to estimated policies and continuous or perceptual state spaces remains unresolved.