Expediting Reinforcement Learning by Incorporating Knowledge About Temporal Causality in the Environment

TL;DR

The paper tackles the difficulty of injecting temporal causal knowledge into reinforcement learning when rewards are sparse and depend on sequences of events. It proposes merging Probabilistic Reward Machines (PRMs) with Temporal Logic-based Causal Diagrams (TL-CDs) to encode temporal causality as a causal DFA, and then forming a product PRM to enable look-ahead during learning. A robust method constructs intermediate products $\mathsf{B_1} = \mathcal{C} \times \mathsf{A}$ and $\mathsf{B_2} = \mathcal{C} \times (-\mathsf{A})$, introduces a minimal reward $m$ for sink transitions, and computes state values via $v^{\star}(u)$ to derive a final PRM $\mathsf{B}$ that preserves the original optimal policy while accelerating convergence. Theoretical convergence to the optimal policy is established (Theorem), and empirical results on coffee-vs-soda and other gridworld tasks demonstrate substantial gains in sample efficiency and robustness to noisy causal inputs. This approach enables more efficient transfer of task specifications across environments by leveraging temporal causality in the reward structure.

Abstract

Reinforcement learning (RL) algorithms struggle with learning optimal policies for tasks where reward feedback is sparse and depends on a complex sequence of events in the environment. Probabilistic reward machines (PRMs) are finite-state formalisms that can capture temporal dependencies in the reward signal, along with nondeterministic task outcomes. While special RL algorithms can exploit this finite-state structure to expedite learning, PRMs remain difficult to modify and design by hand. This hinders the already difficult tasks of utilizing high-level causal knowledge about the environment, and transferring the reward formalism into a new domain with a different causal structure. This paper proposes a novel method to incorporate causal information in the form of Temporal Logic-based Causal Diagrams into the reward formalism, thereby expediting policy learning and aiding the transfer of task specifications to new environments. Furthermore, we provide a theoretical result about convergence to optimal policy for our method, and demonstrate its strengths empirically.

Figures (11)

• Figure 1: An MDP (left) and a PRM (right) that captures the task of bringing either coffee or soda to the office. The coffee machine has a probability of $10\%$ to malfunction and produce bad coffee, leading to a reduced reward of $0.1$ instead of $1$. Bringing soda to the office results in a reward of $1$ deterministically. An example input for the PRM is $\{ \texttt{c, s} \}, \emptyset, \{ \texttt{o, c} \}$ (a sequence of three labels), which will induce the run $q_0 \mapsto q_3 \mapsto q_3 \mapsto q_4$ with a reward of $1$. It is important to note that inputs for PRMs are sets of descriptive propositional variables that are true in a given step, hence why a single label such as $\{ \texttt{c, s} \}$ can include multiple (or $0$) variables.
• Figure 2: Figure \ref{['fig:cd-soda']} (left) is the TL-CD which captures relevant causal information in the environment from Figure \ref{['fig:gridworld-coffee-soda']}. Figure \ref{['fig:cd-collection']} (right) is a TL-CD that holds for the case study in Figure \ref{['fig:mdp-collection_second_case']}.
• Figure 3: Two factors of the causal DFA for the TL-CD in Figure \ref{['fig:cd-soda']}. Rejecting sink states are diamond-shaped. Their parallel composition is the true causal DFA, and its states come from the Cartesian product of states in this Figure. For example, the initial state is $(u_0, t_0)$.
• Figure 4: A fragment of the product of the PRM from Figure \ref{['fig:prm-coffee-soda']} and the TL-CD from Figure \ref{['fig:cd-soda']}. Inheriting the $q$, $u$, and $t$ names of PRM and causal DFA states from previous figures, $x_0 = (q_0, u_0, t_0)$, $x_1 = (q_3, u_1, t_0)$, $x_2 = (q_4, u_3, t_0)$, $x_3 = (q_3, u_2, t_1)$, and $x_4 = (q_4, u_2, t_2)$. Due to the maximum in Equation \ref{['eqn:state-value']}, dashed transitions do not contribute to state value. Dashed states $x_1$ and $x_3$ have $0$ value in both $B_1$ (depicted) and $B_2$, and will be added to the set of terminal states.
• Figure 5: Reward per step averaged over $20$ runs. "No causal" refers to using QRM with the original PRM that does not account for additional causal information in the environment. "Causal" are the results for our method. Both graphs showcase QRM convergence to the optimal policy.

Theorems & Definitions (13)

• definition thmcounterdefinition: Probabilistic Reward Machine (PRM)
• definition thmcounterdefinition: Deterministic Finite Automaton (DFA)
• definition thmcounterdefinition: PRM & TL-CD product
• Theorem 1: Convergence to Optimal Policy
• Lemma 1: Transformation 1
• proof
• definition thmcounterdefinition: Unreachable PRM state
• Lemma 2: Transformation 2
• proof
• definition thmcounterdefinition: Dependent Set of Unreachable PRM States