Reinforcement Learning with Stochastic Reward Machines

AI summary

Overview

• Research area: Reinforcement learning, specifically non-Markovian reward learning, reward machines, and constraint/SMT-based program synthesis.
• Technical level: Advanced. Familiarity with Markov decision processes, reward machines, Q-learning, and constraint solving is assumed.
• Scope: The paper introduces stochastic reward machines and an inference algorithm (SRMI) that learns them from noisy rewards, guaranteeing convergence to an optimal policy in the limit.

What This Paper Is About

Reinforcement learning agents often need rewards that depend on long sequences of actions rather than just the current state. Reward machines solve this by representing such non-Markovian rewards as finite automata, but existing methods assume rewards are noise-free. This paper introduces stochastic reward machines, in which each transition outputs a probability distribution over rewards rather than a fixed number, and presents an algorithm that learns such machines from an agent's noisy experience.

Key Contributions

1. Definition of stochastic reward machines (SRMs). A generalization of reward machines where transitions emit cumulative distribution functions instead of real numbers, letting the formalism natively represent noisy, non-Markovian rewards. Classical reward machines become a special case where distributions are degenerate.

2. A notion of equivalence in expectation. The paper proves that two SRMs need only agree on expected reward values (not exact distributions) to induce the same optimal policy, which substantially simplifies inference and reduces the information the learner must recover.

3. SRMI: a constraint-based inference algorithm. The new algorithm encodes εc-consistency between a candidate SRM and the observed traces as a formula in propositional and real arithmetic, extending the JIRP encoding of Xu et al. It searches for the smallest consistent SRM and uses an Estimates step to refine output means from empirical rewards.

4. Convergence and empirical validation. The paper proves SRMI converges in the limit to an SRM equivalent in expectation to the true one, and therefore to an optimal policy, and demonstrates this on two noisy non-Markovian case studies, outperforming existing RM inference methods and a sample-averaging baseline.

Main Findings

• Classical reward machines cannot handle noisy rewards: Because they output single real values, they either find no consistent machine when noise is present or overfit and blow up in size. This is a fundamental limitation, not a tuning issue.

• Stochastic reward machines capture noise natively: Transitions are annotated with bounded continuous distributions (e.g., uniform intervals), so a label sequence produces a sequence of distributions that are sampled to yield rewards. This preserves interpretability while modeling reward uncertainty.

• Equivalence in expectation is sufficient: Lemma 1 shows that SRMs matching only expected outputs induce identical optimal policies, so the inference problem can focus on recovering means rather than full distributions.

• A sample-replay baseline is brittle and expensive: The obvious fix—replaying trajectories many times to average out noise before inference—requires on-demand repeated sampling, which is costly and sometimes impossible in real environments.

• SRMI avoids replay by solving constraints directly: Counterexamples are handled immediately: Type 1 counterexamples (fixable by shifting outputs) are repaired without structural change, while Type 2 counterexamples trigger a constraint solve for a new minimal SRM.

• The Estimates step resolves non-uniqueness: Many SRMs are εc-consistent with the same trace set, so the algorithm simulates the inferred machine on consistent traces and replaces outputs with midrange estimates from empirical rewards, guaranteeing convergence to the correct expectations.

• Empirical results on two case studies: SRMI beats existing (noise-ignoring) RM inference algorithms and the averaging baseline, confirming that explicitly modeling stochasticity pays off in noisy non-Markovian settings.

• Theoretical guarantee: Theorem 4.1 establishes convergence in the limit under Assumption 1 (roughly, that noise from one distribution cannot fully mask a differently-meaned distribution), and Corollary 1 extends this to convergence to an optimal policy.

Methodology in Plain English

The researchers generalize reward machines so that each transition emits a probability distribution over rewards instead of a fixed number. To learn such a machine, an agent runs a reward-machine-aware Q-learning variant (QRM), which treats the machine as extra memory and uses sampled rewards in place of deterministic ones.

As the agent explores, it collects traces—sequences of observed events paired with rewards. When a trace contradicts the current hypothesis machine (a "counterexample"), the algorithm either patches the current machine by shifting its output means, or, if that is insufficient, solves a constraint problem to synthesize a brand-new, minimal machine consistent with all counterexamples so far. The constraint encoding uses Boolean variables for transitions and real variables for expected rewards, and requires each observed reward to fall within a known noise tolerance of the corresponding expected value.

Once a machine is found, the algorithm consults all consistent traces collected so far and recomputes the output means using midrange estimates, filtering out traces that do not fit. This final refinement step is what guarantees the hypothesis sequence converges to the true reward structure rather than settling on an arbitrary consistent alternative. The whole loop repeats, gradually shrinking the space of plausible machines.

Why This Matters

Impact on research. The paper closes a practical gap between reward machine theory and real reinforcement learning, where rewards are essentially never noise-free. It extends constraint-based RM inference (JIRP) into the stochastic regime, introduces a formal equivalence relation suited to noisy rewards, and provides convergence guarantees. This creates a foundation for subsequent work on robust, noise-tolerant non-Markovian RL.

Real-world applications:

• Robotics and manufacturing: Tasks like assembling a product in the correct order, where sensor readings and material quality vary, require noisy sequential rewards.

• Finance and trading: Rewards from market transactions depend on complex multi-step strategies and are inherently stochastic due to price fluctuations.

• Game playing and simulation: Environments like Minecraft-style resource gathering feature variable payoffs from item quality or randomized market prices.

• Healthcare and process control: Multi-stage treatment or industrial procedures where outcomes are probabilistic but the correct ordering of steps matters.

Industry relevance. Any deployed RL system with delayed, sequence-dependent, and noisy feedback stands to benefit. The algorithm's ability to produce interpretable machines that explicitly capture reward distributions is especially valuable in regulated settings where explaining an agent's objective matters as much as its performance.

Future Directions

• Relaxing Assumption 1. The convergence guarantee requires that noise from one output distribution cannot fully conceal the signal of a differently-meaned one. Characterizing or weakening this condition would broaden applicability.

• Extending beyond symmetric distributions. The main text assumes symmetric noise for simplicity; the appendix sketches an extension, but scaling asymmetric and multimodal distributions to large problems remains open.

• Scaling constraint solving. The current approach enumerates machine sizes and solves increasingly large real-arithmetic formulas, which may bottleneck on environments with many propositions or long-horizon dependencies.

• Bridging to richer noise models. Related literature (Everitt et al., Romoff et al., Wang et al.) handles adversarial or reward-flipping noise; combining those models with the SRM formalism could yield more robust learners.

Target Audience

Researchers and graduate students in reinforcement learning, formal methods, and AI safety who work on non-Markovian rewards, reward machines, or program synthesis from interaction data. It is also relevant to practitioners building RL systems where feedback is delayed, sequence-dependent, and noisy, who need interpretable and provably convergent reward representations. Readers should be comfortable with MDPs, automata, Q-learning, and basic constraint/SMT solving.

Authors’ abstract

Reward machines are an established tool for dealing with reinforcement learning problems in which rewards are sparse and depend on complex sequences of actions. However, existing algorithms for learning reward machines assume an overly idealized setting where rewards have to be free of noise. To overcome this practical limitation, we introduce a novel type of reward machines, called stochastic reward machines, and an algorithm for learning them. Our algorithm, based on constraint solving, learns minimal stochastic reward machines from the explorations of a reinforcement learning agent. This algorithm can easily be paired with existing reinforcement learning algorithms for reward machines and guarantees to converge to an optimal policy in the limit. We demonstrate the effectiveness of our algorithm in two case studies and show that it outperforms both existing methods and a naive approach for handling noisy reward functions.