Learning Task Automata for Reinforcement Learning using Hidden Markov Models

TL;DR

This paper introduces a method to learn task automata from agent experiences in unknown environments, improving reinforcement learning efficiency and interpretability by modeling non-Markovian rewards with hidden Markov models.

Contribution

The paper presents a novel pipeline that learns deterministic finite automata representing tasks from episodes, using hidden Markov models and product MDPs, enhancing transferability and interpretability.

Findings

Effective learning of task automata from limited episodes

Improved policy synthesis speed due to task decomposition

Automata are environment-agnostic and suitable for transfer learning

Abstract

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.