[Paper Review] Joint Inference of Reward Machines and Policies for Reinforcement Learning

Zhe Xu, Ivan Gavran+5 more|arXiv (Cornell University)|Sep 12, 2019

Receptor Mechanisms and Signaling12 citations

TL;DR

This paper proposes Joint Inference of Reward Machines and Policies (JIRP), an iterative algorithm that jointly learns reward machines and policies in reinforcement learning when high-level reward structure is initially unknown. By iteratively refining hypotheses using counterexamples from RL episodes and transferring Q-functions across reward machine states, JIRP converges almost surely to an optimal policy and outperforms standard Q-learning and hierarchical RL in sparse-reward environments like autonomous driving, office navigation, and Minecraft tasks.

ABSTRACT

Incorporating high-level knowledge is an effective way to expedite reinforcement learning (RL), especially for complex tasks with sparse rewards. We investigate an RL problem where the high-level knowledge is in the form of reward machines, i.e., a type of Mealy machine that encodes the reward functions. We focus on a setting in which this knowledge is a priori not available to the learning agent. We develop an iterative algorithm that performs joint inference of reward machines and policies for RL (more specifically, q-learning). In each iteration, the algorithm maintains a hypothesis reward machine and a sample of RL episodes. It derives q-functions from the current hypothesis reward machine, and performs RL to update the q-functions. While performing RL, the algorithm updates the sample by adding RL episodes along which the obtained rewards are inconsistent with the rewards based on the current hypothesis reward machine. In the next iteration, the algorithm infers a new hypothesis reward machine from the updated sample. Based on an equivalence relationship we defined between states of reward machines, we transfer the q-functions between the hypothesis reward machines in consecutive iterations. We prove that the proposed algorithm converges almost surely to an optimal policy in the limit if a minimal reward machine can be inferred and the maximal length of each RL episode is sufficiently long. The experiments show that learning high-level knowledge in the form of reward machines can lead to fast convergence to optimal policies in RL, while standard RL methods such as q-learning and hierarchical RL methods fail to converge to optimal policies after a substantial number of training steps in many tasks.

Motivation & Objective

• To address the challenge of sparse-reward reinforcement learning where high-level task structure is not initially available to the agent.
• To develop a method that jointly infers reward machines and policies in an iterative manner, improving sample efficiency and convergence speed.
• To enable the agent to discover implicit task hierarchies and structural relationships among subtasks through experience and counterexamples.
• To prove almost-sure convergence to an optimal policy under minimal reward machine inference and sufficient episode length.
• To outperform standard Q-learning and hierarchical RL in complex, sparse-reward environments such as autonomous navigation and Minecraft tasks.

Proposed method

• Maintains a hypothesis reward machine and a sample of RL episodes, iteratively refining both using counterexamples where observed rewards contradict the current hypothesis.
• Derives Q-functions from the current reward machine hypothesis and performs standard Q-learning updates using these Q-functions.
• Collects and adds counterexample episodes—where actual rewards differ from those predicted by the current reward machine—into the sample to improve hypothesis inference.
• Uses automata learning techniques to infer a new, more accurate reward machine hypothesis from the updated sample of episodes.
• Transfers Q-functions between equivalent states of consecutive reward machine hypotheses to accelerate learning and improve sample efficiency.
• Employs a polynomial-time learning algorithm for efficient reward machine inference and a periodic batch update strategy to control inference frequency.

Experimental results

Research questions

• RQ1 Can a reinforcement learning agent jointly infer a reward machine and an optimal policy when the high-level task structure is initially unknown?
• RQ2 How can counterexamples from RL episodes be effectively used to refine reward machine hypotheses over time?
• RQ3 What conditions ensure almost-sure convergence to an optimal policy in the joint inference framework?
• RQ4 How does the transfer of Q-functions between equivalent states of successive reward machines improve learning efficiency?
• RQ5 To what extent does the proposed method outperform standard Q-learning and hierarchical RL in sparse-reward environments?

Key findings

• In the autonomous vehicle scenario, JIRP converged to an optimal policy within 100,000 training steps, while baseline methods remained stuck with near-zero cumulative reward after two million steps.
• In both the office world and Minecraft scenarios, JIRP achieved optimal average cumulative reward, whereas baseline methods reached only about 60% of the optimal performance over the same training duration.
• The inferred reward machines in the final iterations accurately captured the required task orderings and dependencies, such as 'get wood before using the workbench' or 'use factory after gathering materials'.
• The use of counterexample collection significantly improved hypothesis quality, with reward machines in later iterations correctly modeling task structure and reward logic.
• Q-function transfer between equivalent states of successive reward machines contributed to faster convergence and reduced sample complexity.
• The algorithm demonstrated robustness across diverse environments, including navigation, resource gathering, and crafting tasks, with consistent performance gains over baselines.