# Learning Reward Machines from Partially Observed Policies

arXiv: https://arxiv.org/abs/2502.03762 · HTML: https://mlsanity.com/paper/2502.03762/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Inverse reinforcement learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 17 of its 48 references that are in the map.
- First posted on arXiv 6 Feb 2025.

## Placed in

- Offline RL and imitation learning from data / [Inverse reinforcement learning](https://mlsanity.com/offlinerl/inverse-reinforcement-learning/)

## Builds on (most cited first, 17 of 17)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [LTL2Action: Generalizing LTL Instructions for Multi-Task RL](https://mlsanity.com/paper/2102.06858/) (2021)
- [Joint Inference of Reward Machines and Policies for Reinforcement Learning](https://mlsanity.com/paper/1909.05912/) (2019)
- [Identifiability in inverse reinforcement learning](https://mlsanity.com/paper/2106.03498/) (2021)
- [Learning Linear Temporal Properties](https://mlsanity.com/paper/1806.03953/) (2018)
- [The Logical Options Framework](https://mlsanity.com/paper/2102.12571/) (2021)
- [Active Finite Reward Automaton Inference and Reinforcement Learning Using Queries and Counterexamples](https://mlsanity.com/paper/2006.15714/) (2020)
- [Induction of Subgoal Automata for Reinforcement Learning](https://mlsanity.com/paper/1911.13152/) (2019)
- [Active Task-Inference-Guided Deep Inverse Reinforcement Learning](https://mlsanity.com/paper/2001.09227/) (2020)
- [Learning Reward Machines: A Study in Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2112.09477/) (2021)
- [Explaining Multi-stage Tasks by Learning Temporal Logic Formulas from Suboptimal Demonstrations](https://mlsanity.com/paper/2006.02411/) (2020)
- [Learning to Search in Task and Motion Planning with Streams](https://mlsanity.com/paper/2111.13144/) (2021)
- [Multi-intention Inverse Q-learning for Interpretable Behavior Representation](https://mlsanity.com/paper/2311.13870/) (2023)
- [Learning Task Automata for Reinforcement Learning using Hidden Markov Models](https://mlsanity.com/paper/2208.11838/) (2022)
- [Learning Robust Reward Machines from Noisy Labels](https://mlsanity.com/paper/2408.14871/) (2024)
- [Learning Moore Machines from Input-Output Traces](https://mlsanity.com/paper/1605.07805/) (2016)
- [Reward Machine Inference for Robotic Manipulation](https://mlsanity.com/paper/2412.10096/) (2024)

## Built on by (most cited first, 3 of 3)

- [Average Reward Reinforcement Learning for Omega-Regular and Mean-Payoff Objectives](https://mlsanity.com/paper/2505.15693/) (2025)
- [Fitting Reinforcement Learning Model to Behavioral Data under Bandits](https://mlsanity.com/paper/2511.04454/) (2025)
- [Active Reward Machine Inference From Raw State Trajectories](https://mlsanity.com/paper/2604.07480/) (2026)

## Abstract (arXiv)

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.
