Pushdown Reward Machines for Reinforcement Learning

TL;DR

This work extends reward machines by introducing pushdown reward machines (pdRMs) that leverage a single stack to represent deterministic context-free languages, enabling rewards for temporally extended behaviors beyond regular languages. It defines two policy families for pdRMs—full-stack policies and top-$k$ policies—and provides a criterion to determine when top-$k$ policies can attain the same value as full policies, along with formal analyses of expressive power and space complexity. The paper compares pdRMs with counting reward automata (CRAs), showing that pdRMs can offer more compact representations in certain regimes while acknowledging CRA expressiveness can exceed that of pdRMs for multiple counters. It also introduces CpRM, a counterfactual learning extension for pdRMs, and demonstrates through five domains that pdRMs can achieve strong sample efficiency and competitive performance in both discrete and continuous settings, sometimes outperforming recurrent neural approaches. Overall, pdRMs provide a principled, scalable framework for encoding and learning with deterministic context-free rewards, with practical benefits in tasks requiring memory of structured sequences.

Abstract

Reward machines (RMs) are automata structures that encode (non-Markovian) reward functions for reinforcement learning (RL). RMs can reward any behaviour representable in regular languages and, when paired with RL algorithms that exploit RM structure, have been shown to significantly improve sample efficiency in many domains. In this work, we present pushdown reward machines (pdRMs), an extension of reward machines based on deterministic pushdown automata. pdRMs can recognise and reward temporally extended behaviours representable in deterministic context-free languages, making them more expressive than reward machines. We introduce two variants of pdRM-based policies, one which has access to the entire stack of the pdRM, and one which can only access the top $k$ symbols (for a given constant $k$) of the stack. We propose a procedure to check when the two kinds of policies (for a given environment, pdRM, and constant $k$) achieve the same optimal state values. We then provide theoretical results establishing the expressive power of pdRMs, and space complexity results for the proposed learning problems. Lastly, we propose an approach for off-policy RL algorithms that exploits counterfactual experiences with pdRMs. We conclude by providing experimental results showing how agents can be trained to perform tasks representable in deterministic context-free languages using pdRMs.

Figures (17)

• Figure 1: Pushdown reward machine for the Maze task. Each transition is labelled with a tuple $\ell, z / \zeta, r$, where $\ell$ is the input observation, $z$ the top symbol on the stack, $\zeta$ the string of symbols pushed onto the stack (with the new top symbol leftmost in $\zeta$), and $r$ is the output reward. The symbol "$z$" indicates an arbitrary symbol in $\Gamma$. In the transition from $u_1$ to $u_2$, $a'$ represents any direction except for $\overline{a}$, the opposite direction to $a$.
• Figure 2: LetterEnv results, comparing agents trained with a pdRM and with a CRA. This shows how pdRMs can be used to encode part of the tasks encodable by CRAs.
• Figure 3: 1-TreasureMaze (left) and MultipleTreasureMaze (right) results. We provide individual plots for each maze in \ref{['sec:Appendix-E']} of the Appendix. In both plots, each maze is identified by a line style, and each agent by a colour. The 1-TreasureMaze plot shows how the 1-pdRM agent could achieve the task on all mazes, whereas the pdRM and the Q-learning CRA agents only on the smallest maze. The -CpRM agents only managed to achieve the task in the smallest maze; in the other cases training timed out due to the time required to perform the policy updates. The CQL-CRA agent never managed to achieve the task. For the MultipleTreasureMaze experiment, we include only results from the 1-pdRM and pdRM agents trained with vanilla Q-learning. As can be seen, on mazes the 1-pdRM agent learnt to achieve the task, whereas the pdRM agent did not. We believe this is due to the smaller size of the $1$-policy.
• Figure 4: DeliverWorld results. Left plot: agents performed 8 deliveries during training and testing (DeliverWorld-8). Right plot: agents performed 4 deliveries during training episodes and 8 during testing episodes (DeliverWorld-4-8). For DeliverWorld-8, the only agents that consistently achieve the task at the end of training are the -CpRM agents, the Q-learning 1-pdRM agent and the hierarchical 1-pdRM and 2-pdRM agents. The hierarchical pdRM agent manages to achieve the task but not as consistently. The other Q-learning agents show the worse performance out of all agents, with the pdRM agent flatlining at a reward of $-1$. On the other hand, for DeliverWorld-4-8, only the -CpRM agents and the hierarchical agents eventually learn to consistently achieve the task. This shows how pdRMs can help in obtaining agents that complete longer tasks than the ones they were trained in.
• Figure 5: WaterWorld results, 17$\times$17 map (left) and 20$\times$20 map (right). We compare the performance of a 1-pdRM agent trained with PPO against that of an agent trained with PPO and one of an agent trained with recurrent PPO. In the 17$\times$17 map, the 1-pdRM agent is able to achieve the task very quickly. Only the recurrent PPO agent manages to considerably improve its performance, but does not match that of the 1-pdRM agent. Similarly to the results of the 17$\times$17 map, in the 20$\times$20 map the 1-pdRM agent is the only one that consistently achieves the task; significanly outperforming the other agents which do not improve their performance. These two plots show how the pdRM is crucial in training agents to achieve this task.

Theorems & Definitions (11)

• Definition 1: Pushdown Reward Machine
• Definition 2
• Definition 3: Policy
• Definition 4: $k$-policy
• Definition 5: $k$-equivalence $\sim_k$
• Proposition 1
• proof
• Theorem 1
• Theorem 2
• Theorem 3