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

2025-10-17

AI Summary

Existing probabilistic reward machines (PRMs) struggle to incorporate high-level causal knowledge for sparse-reward reinforcement learning tasks requiring complex event sequences, limiting interpretability and cross-domain transferability of reward design. Method: This paper introduces the Temporal Logic Causal Graph (TLCG) into reward modeling, enabling an interpretable and editable structured reward mechanism. We propose the TLCG-PRM framework, which semantically embeds causal graph structure into PRMs and jointly optimizes them with standard RL algorithms. Contribution/Results: We provide theoretical guarantees that the method converges to the optimal policy. Empirical evaluation across multiple complex sequential decision-making tasks demonstrates a 2.3× average acceleration in learning convergence. Moreover, the causal structure enables effective reward transfer across environments—significantly improving generalization without retraining reward models from scratch.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Knowledge Representation and Reasoning: Action, Change, and Causality

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

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.

Problem

• Expediting reinforcement learning for sparse reward tasks
• Incorporating temporal causality knowledge into reward formalism
• Enabling transfer of task specifications across environments

Innovation

• Incorporating Temporal Logic-based Causal Diagrams
• Expediting policy learning with causal knowledge
• Transferring task specifications to new environments