Efficient Reinforcement Learning in Probabilistic Reward Machines

In this paper, we study reinforcement learning in Markov Decision Processes with Probabilistic Reward Machines (PRMs), a form of non-Markovian reward commonly found in robotics tasks. We design an algorithm for PRMs that achieves a regret bound of Õ((HOAT)^(1/2) + H²O²A^(3/2) + H(T)^(1/2)), where H is the time horizon, O is the number of observations, A is the number of actions, and T is the number of time steps. This result improves over the best-known bound, Õ(H(OAT)^(1/2)), for MDPs with Deterministic Reward Machines (DRMs), a special case of PRMs. When T ≥ H³O³A² and OA ≥ H, our regret bound leads to a regret of Õ((HOAT)^(1/2)), which matches the established lower bound of Ω((HOAT)^(1/2)) for MDPs with DRMs up to a logarithmic factor. To the best of our knowledge, this is the first efficient algorithm for PRMs. Additionally, we present a new simulation lemma for non-Markovian rewards, which enables reward-free exploration for any non-Markovian reward given access to an approximate planner. Complementing our theoretical findings, we show through extensive experimental evaluations that our algorithm indeed outperforms prior methods in various PRM environments.