Model-Based Reinforcement Learning in Discrete-Action Non-Markovian Reward Decision Processes

AI summary

Overview

Research area: Reinforcement learning — specifically model-based RL for non-Markovian reward decision processes (NMRDPs), reward machines, and continuous-state discretisation.

Technical level: Advanced. The paper combines algorithm design with PAC-MDP sample-complexity analysis (lemmas, thresholds, union bounds) and assumes familiarity with MDPs, value iteration, and probably-approximately-correct learning theory.

Scope (one sentence): The paper introduces QR-Max, a model-based algorithm that factorises Markovian environment dynamics from non-Markovian reward-machine dynamics to obtain PAC guarantees and improved sample efficiency in discrete-action NMRDPs, plus Bucket-QR-Max, a SimHash-based extension to continuous states.

What This Paper Is About

Many tasks succeed or fail based on the whole history of interaction, not just the current state — so a Markovian reward gives a brittle or misleading learning signal. Such tasks are typically modelled as NMRDPs and then reduced to an ordinary MDP over the product space S × Q (environment state × automaton state). The paper's goal is to show that this product-space view wastes effort, because the environment transitions are still Markovian and identical across automaton states, and that explicitly separating the two dynamics yields a provably more sample-efficient algorithm.

Key Contributions

1. QR-Max, a model-based RL algorithm for discrete-action NMRDPs that factorises learning into a Markovian environment component P(s′|s,a), R_E(s,a,s′) and a non-Markovian automaton component P(q′|s′,q), R_A(q,s′,q′), each with its own knownness threshold t_E and t_Q. Each learned environment transition is reused across every automaton state.
2. A PAC-MDP guarantee with a factorised sample bound. The unknown-component count is U = |S||A|m_E + |S||Q|m_Q, replacing the product-space count |S||Q||A|m incurred by applying R-Max directly on S × Q. The paper states this is the first model-based RL algorithm for discrete-action NMRDPs that exploits this decoupling to obtain PAC convergence to ε-optimal policies with polynomial sample complexity.
3. Bucket-QR-Max, a continuous-state extension using SimHash (Charikar 2002) locality-sensitive hashing to build buckets on the fly, with no manual gridding and no function approximation. It preserves the factorised counters and optimism, and yields sample complexity in |B||A| + |B||Q| rather than the cubic |B||Q||A|. It admits a conditional PAC-MDP bound under a bounded-discretisation assumption (Assumption B.1, Theorem B.2).
4. An empirical comparison against UCBVI-sB, UCBVI-H, UCBVI-B, PSRL, OPSRL, R-Max, and the model-free QRM on environments of increasing complexity, including a continuous-state variant.

Main Findings

• Large sample-efficiency gains: On the benchmarks considered, QR-Max is reported as one to two orders of magnitude more sample-efficient than standard R-Max, reward-machine-aware QRM, and optimistic explorers UCBVI, PSRL and OPSRL, while retaining formal PAC guarantees. The authors also report greater robustness in finding optimal policies.
• Theoretical saving comes from the automaton side being cheap: Since the automaton is deterministic, t_Q = 1 suffices to know all of its states (m_Q = 1 and ξ_Q = 0 for deterministic reward machines with deterministic transition rewards). The dominant learned component is therefore the shared environment model, which is where the advantage over product-space R-Max arises.
• Memory matters on large grids: On the 15×15 grid with H = 250, PSRL/OPSRL would need approximately 50.6M entries versus 0.204M for QR-Max, exceeding the authors' budget — so PSRL/OPSRL are omitted when the state space is very large. UCBVI has memory O(S²A) and PSRL/OPSRL have memory O(H, S²A), a quadratic dependence on S.
• PAC-MDP bound (Theorem 3): With probability at least 1 − δ, the number of interaction steps in which the greedy policy is not ε-optimal is bounded by O((HU/p) log(U/δ)), where H = ⌈(1−γ)⁻¹ ln(2R_max/[ε(1−γ)])⌉, p = ε(1−γ)/(2R_max), and U = |S||A|m_E + |S||Q|m_Q.
• Two supporting lemmas: Lemma 1 (Factored Accuracy) bounds reward error by ξ_R and total-variation transition error by ξ_E + ξ_Q on every known product component; Lemma 2 (Near-Optimality in the Known Model) shows the greedy policy is ε-optimal whenever no unknown component is reached within the effective horizon.
• RM knowledge is optional: QR-Max does not require the automaton transition function η as input — it only needs the observed automaton-state signal q and q′. When the reward machine is given, the structure-aware variant QR-MaxRM additionally generates counterfactual experience across automaton states.
• Continuous extension is nearly the same algorithm: Bucket-QR-Max differs from discrete QR-Max in exactly three lines — LSH discretisation b ← h(s_cont), four separate counters (n_{E_T}, n_{E_R}, n_{Q_T}, n_{Q_R}) each with its own threshold, and bucket-indexed rather than state-indexed tables. All four conditions must be met before (b,a,q) is marked known and value iteration is triggered.
• Guarantee caveat for hashing: SimHash does not deterministically bound bucket diameters, so the PAC guarantee is conditional on the induced partition satisfying the bounded-diameter assumption, with hash collisions absorbed into an extra failure probability δ_LSH (overall success at least 1 − (δ + δ_LSH)).
• Not directly comparable to some prior work: Regret-based methods for deterministic RMs and probabilistic RMs optimise regret (often average-reward or finite-horizon) under a known RM/PRM and are therefore not directly comparable to this paper's discounted PAC-style evaluation. The paper also notes Shao and Kwiatkowska (2023) is only indirectly compared, since on the problems considered it corresponds to QRM.
• Factorisation is orthogonal to factored MDPs: Here the underlying environment need not be factored; the gain comes purely from separating Markovian environment dynamics from the non-Markovian reward automaton.

Methodology in Plain English

The starting point is the standard trick for non-Markovian rewards: attach a finite automaton (such as a reward machine or restraining bolt) to the environment, and let the reward depend on the automaton state q together with the observed successor state s′. Running R-Max on the combined (s, q) state works in theory but relearns the same environment transition once for every automaton state.

The authors instead write the joint transition as a product, P(s′,q′|s,q,a) = P(s′,a)P(q′|s′,q), and the reward as a sum R_E(s,a,s′) + R_A(q,s′,q′). This holds because the environment is Markovian (so s′ does not depend on q) and because the automaton is deterministic (so q′ depends only on q and s′).

QR-Max then keeps two separate sets of counters: how often each environment (s,a) pair has been tried, and how often each automaton (q,s′) pair has been seen. The Q-table is initialised optimistically at R_max/(1−γ) for all state-action pairs. Actions are chosen greedily from the Q-table. When an occurrence count reaches its threshold, the corresponding component is labelled "known", and value iteration is rerun over the model, using only known components and leaving everything else at its optimistic initial value — the standard R-Max optimism that drives exploration. Estimators are simple visit ratios: empirical counts for the transition probabilities and accumulated reward divided by visit counts for the rewards, expressed in "reward-mass" form so the backup does not have to condition on rare successor events.

For continuous states, the same machinery runs on buckets produced by L independent SimHash functions, each projecting the continuous state onto {0, …, 2^{d_h} − 1}. Buckets are discovered on the fly and assigned integer indices. Because the hash granularity is fixed a priori, each bucket plays exactly the role a discrete state does in QR-Max.

Experiments compare QR-Max and QR-MaxRM against UCBVI-sB, UCBVI-H, UCBVI-B, PSRL, OPSRL, R-Max, and QRM using the authors' recommended hyper-parameters. Note: the supplied paper content is truncated inside the experiments section (it cuts off mid-sentence at "For finite-horizon baselines that require a"), so the exact experimental results tables, the full list of environments, episode counts, and hyper-parameter values are not reported in the available text. Code is released at https://github.com/Alee08/qrmax.

Why This Matters

Impact on research. The paper reframes a well-known reduction — collapsing an NMRDP into a product MDP — as a source of avoidable exploration cost, and shows the fix comes with formal guarantees rather than just empirical speedups. It also connects model-based RL theory (R-Max style PAC-MDP analysis, Strehl et al. 2009) to the reward-machine literature, which has largely been model-free or regret-based. It positions factorisation as orthogonal to factored MDPs and as distinct from Covering-RMAX, and it explicitly leaves learning the NMR specification itself to prior and future work.

Potential real-world applications (the paper itself motivates the setting through "game playing to robotics" and tasks spanning temporally ordered sub-goals; the specific deployments below are implied rather than reported):

• Robotic task sequences with ordered sub-goals, where success depends on the order of actions rather than a single final state.
• Any domain where a finite-trace temporal specification (LTL_f or LDL_f) encodes the objective and is known in advance.
• Control problems with continuous sensor readings and a small discrete action set, where Bucket-QR-Max avoids hand-designed grids and neural approximations.
• Hierarchical RL pipelines: the authors note that optimal discrete-action solutions can drive effective training in continuous-action MDPs (e.g., Cipollone et al. 2025).

Industry relevance. The paper reports one-to-two-orders-of-magnitude reductions in interaction cost, and interaction steps are the expensive resource in simulation-heavy or hardware-bound settings. The memory comparison is equally practical: on the 15×15 grid with H = 250, approaches with quadratic state dependence need roughly 50.6M entries where QR-Max needs 0.204M. The fact that QR-Max does not need the automaton transition function, only the observed automaton state, also lowers the integration bar for systems that can already expose a monitor signal.

Future Directions

1. Partially observable automaton state. The paper assumes full observability of both s and q. If q is not observable and the reward machine is unknown, the problem becomes partially observable; the authors explicitly leave automaton-state inference to future work.
2. Stochastic reward machines. The threshold t_Q is retained "for modularity and generality" precisely so the algorithm can extend to more general machines such as stochastic RMs (Corazza et al. 2022). Testing that extension is a natural next step.
3. Learning the specification itself. The authors assume the NMR specification is given and place work on learning the NMR model (e.g., Gaon and Brafman 2020) outside scope — a clear open question is combining model learning for the specification with QR-Max's factorised planner.
4. Guarantees without the bounded-discretisation assumption. Bucket-QR-Max's PAC bound is conditional on Assumption B.1 and on SimHash collisions being absorbed into δ_LSH. Removing or weakening that condition — or establishing deterministic bucket-diameter bounds — would strengthen the continuous-state claim.

Target Audience

Researchers and graduate students in reinforcement learning theory and reward-machine / temporal-logic RL who care about sample-complexity guarantees. It is also relevant to practitioners building agents for sequential or temporally specified tasks with discrete actions — especially those whose state space is large enough that product-space methods and quadratic-memory explorers like PSRL/OPSRL become infeasible, and who need a continuous-state variant without hand-tuned discretisation or neural function approximation. Readers without a background in PAC-MDP analysis will find Sections 4.2 and the appendices demanding.

Authors’ abstract

Many practical decision-making problems involve tasks whose success depends on the entire system history, rather than on achieving a state with desired properties. Markovian Reinforcement Learning (RL) approaches are not suitable for such tasks, while RL with non-Markovian reward decision processes (NMRDPs) enables agents to tackle temporal-dependency tasks. This approach has long been known to lack formal guarantees on both (near-)optimality and sample efficiency. We contribute to solving both issues with QR-MAX, a novel model-based algorithm for discrete NMRDPs that factorizes Markovian transition learning from non-Markovian reward handling via reward machines. To the best of our knowledge, this is the first model-based RL algorithm for discrete-action NMRDPs that exploits this factorization to obtain PAC convergence to $\varepsilon$-optimal policies with polynomial sample complexity. We then extend QR-MAX to continuous state spaces with Bucket-QR-MAX, a SimHash-based discretiser that preserves the same factorized structure and achieves fast and stable learning without manual gridding or function approximation. We experimentally compare our method with modern state-of-the-art model-based RL approaches on environments of increasing complexity, showing a significant improvement in sample efficiency and increased robustness in finding optimal policies.