ARM-FM: Automated Reward Machines via Foundation Models for Compositional Reinforcement Learning

TL;DR

This work tackles the challenge of reward design in reinforcement learning by introducing ARM-FM, a framework that uses foundation models to automatically generate Language-Aligned Reward Machines (LARMs) from natural language task descriptions. LARMs provide a compositional, interpretable structure for rewards and embed language descriptions of subtasks into a shared semantic space, enabling knowledge transfer and zero-shot generalization when training a single policy conditioned on language embeddings. The approach yields dense learning signals, improves sample efficiency, and scales to diverse domains including 2D grids, 3D procedurally generated worlds, and robotic manipulation, with strong performance in multi-task settings. Overall, ARM-FM demonstrates that grounding reward design in language-aligned automata can unlock compositional RL with transferable skills and interpretable task representations, while acknowledging a human verification step that could be automated further through formal grounding and self-correction.

Abstract

Reinforcement learning (RL) algorithms are highly sensitive to reward function specification, which remains a central challenge limiting their broad applicability. We present ARM-FM: Automated Reward Machines via Foundation Models, a framework for automated, compositional reward design in RL that leverages the high-level reasoning capabilities of foundation models (FMs). Reward machines (RMs) -- an automata-based formalism for reward specification -- are used as the mechanism for RL objective specification, and are automatically constructed via the use of FMs. The structured formalism of RMs yields effective task decompositions, while the use of FMs enables objective specifications in natural language. Concretely, we (i) use FMs to automatically generate RMs from natural language specifications; (ii) associate language embeddings with each RM automata-state to enable generalization across tasks; and (iii) provide empirical evidence of ARM-FM's effectiveness in a diverse suite of challenging environments, including evidence of zero-shot generalization.

Figures (22)

• Figure 1: An overview of our framework (left) and results in a complex sparse-reward environment (right). Reward Machine Generation (top-left): Given a high-level natural language prompt and a visual observation of the environment, a FM automatically generates the formal specification of the Reward Machine, the executable Python code for the labeling functions, and the natural language descriptions for each RM state. RL training (bottom-left): During the RL training loop, the labeling functions evaluate environment observations to update the Reward Machine's state, which provides a dense reward signal $R_t^{\text{RM}}$. The RL agent's policy receives the environment observation along with the embedding $\phi(\cdot)$ of the current RM state's language description, making it aware of its active sub-goal. Empirical results (right): Results in a complex sparse-reward Minecraft-based resource-gathering task from Craftiummalagon2024craftium, where an RL agent is unable to make progress (top), while our agent, guided by an FM-generated LARM, learns to solve the task efficiently (bottom).
• Figure 2: ARM-FM leverage FMs to automatically construct RMs: using the UnlockToUnlock task description from MiniGrid (left), an RM is automatically constructed to solve the task (right).
• Figure 3: A self-improvement loop where a generator and critic FMs iteratively refine LARMs, with optional human verification.
• Figure 4: The three core components generated by our method for the UnlockToUnlock environment: (Left) the RM specification, (Center) the labeling functions that drive the state transitions, and (Right) the instructions and embeddings for each RM state.
• Figure 5: An overview of the diverse environments used to evaluate ARM-FM. (a) MiniGrid environments test long-horizon planning in sparse-reward settings. (b) Craftium, a 3D world inspired by Minecraft, evaluates how ARM-FM scales with increasing task complexity. (c) Meta-World provides a benchmark for robotic manipulation with continuous control. (d) XLand-MiniGrid is used to evaluate multi-task learning and generalization.