FORM: Learning Expressive and Transferable First-Order Logic Reward Machines

Reward machines (RMs) are an effective approach for addressing non-Markovian rewards in reinforcement learning (RL) through finite-state machines. Traditional RMs, which label edges with propositional logic formulae, inherit the limited expressivity of propositional logic. This limitation hinders the learnability and transferability of RMs since complex tasks will require numerous states and edges. To overcome these challenges, we propose First-Order Reward Machines (textttFORMs), which use first-order logic to label edges, resulting in more compact and transferable RMs. We introduce a novel method for learning textttFORMs and a multi-agent formulation for exploiting them and facilitate their transferability, where multiple agents collaboratively learn policies for a shared textttFORM. Our experimental results demonstrate the scalability of textttFORMs with respect to traditional RMs. Specifically, we show that textttFORMs can be effectively learnt for tasks where traditional RM learning approaches fail. We also show significant improvements in learning speed and task transferability thanks to the multi-agent learning framework and the abstraction provided by the first-order language.