Pushdown Reward Machines in RL
- Pushdown reward machines (pdRMs) are formal models that augment classical reward machines with a pushdown stack to encode deterministic context-free reward functions.
- They extend deterministic finite automata to deterministic pushdown automata, enabling reinforcement learning agents to handle temporally extended, non-Markovian tasks.
- Empirical evaluations across domains like LetterEnv and TreasureMaze demonstrate that top-k pdRM policies achieve efficient state-space management and improved learning performance.
A pushdown reward machine (pdRM) is a formal model augmenting the classical reward machine architecture with a pushdown stack, thereby enhancing its expressive power for encoding non-Markovian reward functions in reinforcement learning. By generalizing reward machines (which correspond to deterministic finite automata) to deterministic pushdown automata with an associated reward function, pdRMs can recognize and assign rewards for behaviors described by deterministic context-free languages, a strict superset of regular languages. This increased expressivity enables reinforcement learning agents to solve a significantly broader class of temporally extended tasks while preserving a principled automata-theoretic foundation (Varricchione et al., 9 Aug 2025).
1. Formal Structure of Pushdown Reward Machines
A pushdown reward machine is a tuple
with the following components:
- : finite set of machine states.
- : initial machine state.
- : set of accepting (final) states.
- : finite input alphabet. .
- : finite stack alphabet. .
- : initial stack symbol.
- : deterministic transition function. For 0, 1 denotes moving from state 2 to 3, popping 4 (unless 5), and pushing 6.
- 7: immediate-reward function emitted during a transition.
This automaton can perform silent (8) moves and stack-zero checks (9 or 0), facilitating standard deterministic pushdown automaton transitions such as arbitrary stack pops/pushes and conditional stack operations.
2. Expressivity: From RMs to pdRMs and Beyond
Classical reward machines (RMs) capture only regular reward structures, corresponding to deterministic finite automata (DFAs) and thus limited to regular languages. The addition of a single pushdown stack (in pdRMs) elevates their expressivity to deterministic pushdown automata (DPDAs), i.e., the class of deterministic context-free languages (DCFLs). This is a strict superset of regular languages, enabling pdRMs to reward temporally extended behaviors that are inexpressible by regular automata.
| Model | Automaton Type | Language Class |
|---|---|---|
| RM | DFA | Regular languages |
| pdRM | DPDA | Deterministic context-free languages (DCFL) |
| CRA (1 ctrs) | Full counter | Recursively enumerable |
This increased expressivity remains short of counter reward automata (CRAs) with multiple counters, which can encode any recursively enumerable reward structure but at greater space and computation cost (Varricchione et al., 9 Aug 2025).
3. Policy Classes and State Representations
In reinforcement learning settings, the product MDP induced by combining an MDP with a pdRM leads to an augmented state space
2
Two canonical policy classes over this space are defined:
- Full-stack policy: 3, where the learned policy can access the entire stack.
- Top-4 policy: 5, where only the top 6 stack symbols are observable.
Formally, two product-states 7 and 8 are equivalent for a top-9 policy iff their top-0 stack prefixes match.
A proposition addresses the sufficiency of top-1 policies: if the optimal full-stack value function and action preferences agree across stack strings sharing their top-2 slice, then the optimal policy can be implemented using only these 3 symbols, yielding a large reduction in policy space complexity (Varricchione et al., 9 Aug 2025).
4. Theoretical Properties: Expressiveness and Space Complexity
The expressive hierarchy and corresponding space considerations are as follows:
- Expressive Power (Theorem 1):
- RMs ≡ DFAs ⇒ regular languages.
- pdRMs ≡ DPDAs ⇒ DCFLs (strict superset).
- CRAs (≥2 counters) ⇒ recursively enumerable (strict superset of DCFLs).
- State-space blow-up (Theorem 2): For episode length 4, push width 5, and max push-only 6-steps 7, there are at most
8
reachable stack configurations, demonstrating exponential growth with stack alphabet size and episode horizon.
- Policy complexity (Theorem 3): For a top-9 policy, only 0 stack views are necessary, dramatically reducing the size of policy representations. For 1, this is linear in 2.
- Pushdown vs. Counter Automata (Theorem 4): Any DCFL can be encoded by a 1-counter CRA at exponential cost, and certain languages (e.g., marked palindromes) force a 3-counter machine to 4 configurations, highlighting the succinctness of pushdown stacks for such tasks.
These results collectively establish that pdRMs improve expressivity over RMs while maintaining more tractable state representations than general counter automata (Varricchione et al., 9 Aug 2025).
5. Empirical Evaluation: Domains and Findings
pdRMs were evaluated across five domains, comparing various policy and automaton configurations:
- LetterEnv: Requires counting a sequence of 5s until an event 6, then matching the count with 7s. Top-1-pdRM matches or exceeds the performance of CQL-CRA and outperforms both full-pdRM and Q-CRA, demonstrating competitive sample efficiency with smaller state spaces.
- 1-TreasureMaze: Agents must find treasure then retrace their paths. On large mazes (10x10, 20x20), only top-1-pdRM converges; full-pdRM state space is too large, and CRAs experience counter state explosion.
- MultipleTreasureMaze: Collection and retracing of multiple treasures. Only pdRMs are tractable; top-1-pdRM learns reliably, full-pdRM does not converge within resource constraints.
- DeliverWorld: Sequential item delivery, using hierarchical RL with pdRMs. Hierarchical top-1 and top-2 policies yield order-of-magnitude faster convergence. Top-2 is faster for long sequences; top-1 generalizes better from shorter to longer tasks.
- WaterWorld: Continuous control, requiring phased stack manipulations. Only pdRM-enhanced agents can solve the full task; recurrent agents (PPO, PPO+LSTM) underperform.
A summary of empirical findings is provided below:
| Domain | Key Result | Policy/Automaton |
|---|---|---|
| LetterEnv | Top-1-pdRM ≃ CQL-CRA; smaller state space | Top-1-pdRM, CQL-CRA |
| 1-TreasureMaze | Only top-1-pdRM solves large instances | Top-1-pdRM |
| MultipleTreasure | Only pdRMs tractable; top-1-pdRM learns reliably | Top-1-pdRM |
| DeliverWorld | Hierarchical top-1/top-2 much faster; top-1 generalizes better | Hierarchical pdRM policies |
| WaterWorld | Only pdRM-augmented agent reliably solves task | PPO + top-1-pdRM |
These experiments reinforce the practical advantages of top-8 pdRM policies for balancing expressivity and tractability in RL tasks with context-free structure (Varricchione et al., 9 Aug 2025).
6. Practical Implications and Research Outlook
Pushdown reward machines furnish a principled, automata-theoretic mechanism for specifying non-Markovian reward structures beyond regular languages while maintaining concise and expressive policy representations. Key implications include:
- Strict expressivity gains over RMs, while avoiding the combinatorial state explosion characteristic of multi-counter CRAs.
- For many practical tasks, even a top-1 pdRM suffices for optimal or near-optimal behavior, resulting in product state spaces of size 9.
- Hierarchical policy architectures and counterfactual learning methods promise further improvements in efficiency.
- Symbolic stack representations in pdRMs can outperform recurrent neural baselines in continuous and context-free RL domains.
A plausible implication is that pdRMs offer a scalable and expressive formalism, opening new avenues for specifying and learning temporally extended behaviors in complex reinforcement learning environments, especially where context-free reward dependencies are inherent (Varricchione et al., 9 Aug 2025).