# Learning Task Automata for Reinforcement Learning using Hidden Markov Models

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 60 references that are in the map.
- First posted on arXiv 25 Aug 2022.

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

- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Reward Machines: Exploiting Reward Function Structure in Reinforcement Learning](https://mlsanity.com/paper/2010.03950/) (2020)
- [Benchmarking Model-Based Reinforcement Learning](https://mlsanity.com/paper/1907.02057/) (2019)
- [Environment-Independent Task Specifications via GLTL](https://mlsanity.com/paper/1704.04341/) (2017)
- [Compositional Reinforcement Learning from Logical Specifications](https://mlsanity.com/paper/2106.13906/) (2021)
- [Joint Inference of Reward Machines and Policies for Reinforcement Learning](https://mlsanity.com/paper/1909.05912/) (2019)
- [Logically-Constrained Reinforcement Learning](https://mlsanity.com/paper/1801.08099/) (2018)
- [Reinforcement Learning with Non-Markovian Rewards](https://mlsanity.com/paper/1912.02552/) (2019)
- [Induction and Exploitation of Subgoal Automata for Reinforcement Learning](https://mlsanity.com/paper/2009.03855/) (2020)
- [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)
- [Verifiable and Compositional Reinforcement Learning Systems](https://mlsanity.com/paper/2106.05864/) (2021)
- [Inferring Probabilistic Reward Machines from Non-Markovian Reward Processes for Reinforcement Learning](https://mlsanity.com/paper/2107.04633/) (2021)
- [Joint Learning of Reward Machines and Policies in Environments with Partially Known Semantics](https://mlsanity.com/paper/2204.11833/) (2022)
- [Statistical and Computational Guarantees for the Baum-Welch Algorithm](https://mlsanity.com/paper/1512.08269/) (2015)
- [Learning and Solving Regular Decision Processes](https://mlsanity.com/paper/2003.01008/) (2020)
- [Markov Abstractions for PAC Reinforcement Learning in Non-Markov Decision Processes](https://mlsanity.com/paper/2205.01053/) (2022)
- [Online Learning of Non-Markovian Reward Models](https://mlsanity.com/paper/2009.12600/) (2020)
- [Quantitative model-checking of controlled discrete-time Markov processes](https://mlsanity.com/paper/1407.5449/) (2014)

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

- [On the Limitations of Markovian Rewards to Express Multi-Objective, Risk-Sensitive, and Modal Tasks](https://mlsanity.com/paper/2401.14811/) (2024)
- [Reinforcement Learning with Knowledge Representation and Reasoning: A Brief Survey](https://mlsanity.com/paper/2304.12090/) (2023)
- [Learning Reward Machines from Partially Observed Policies](https://mlsanity.com/paper/2502.03762/) (2025)
- [Reinforcement Learning with Temporal-Logic-Based Causal Diagrams](https://mlsanity.com/paper/2306.13732/) (2023)
- [Provably Efficient Exploration in Reward Machines with Low Regret](https://mlsanity.com/paper/2412.19194/) (2024)
- [Active Reward Machine Inference From Raw State Trajectories](https://mlsanity.com/paper/2604.07480/) (2026)
- [Exact Distinguishability in Non-Markovian Decision Processes](https://mlsanity.com/paper/2610.01527/) (2026)

## Abstract (arXiv)

Training reinforcement learning (RL) agents using scalar reward signals is often infeasible when an environment has sparse and non-Markovian rewards. Moreover, handcrafting these reward functions before training is prone to misspecification, especially when the environment's dynamics are only partially known. This paper proposes a novel pipeline for learning non-Markovian task specifications as succinct finite-state `task automata' from episodes of agent experience within unknown environments. We leverage two key algorithmic insights. First, we learn a product MDP, a model composed of the specification's automaton and the environment's MDP (both initially unknown), by treating the product MDP as a partially observable MDP and using the well-known Baum-Welch algorithm for learning hidden Markov models. Second, we propose a novel method for distilling the task automaton (assumed to be a deterministic finite automaton) from the learnt product MDP. Our learnt task automaton enables the decomposition of a task into its constituent sub-tasks, which improves the rate at which an RL agent can later synthesise an optimal policy. It also provides an interpretable encoding of high-level environmental and task features, so a human can readily verify that the agent has learnt coherent tasks with no misspecifications. In addition, we take steps towards ensuring that the learnt automaton is environment-agnostic, making it well-suited for use in transfer learning. Finally, we provide experimental results compared with two baselines to illustrate our algorithm's performance in different environments and tasks.