# Reward Shaping and Action Masking for Compositional Tasks using Behavior Trees and LLMs

arXiv: https://arxiv.org/abs/2605.05795 · HTML: https://mlsanity.com/paper/2605.05795/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 30 references that are in the map.
- First posted on arXiv 7 May 2026.

## Builds on (most cited first, 21 of 21)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [OpenAI GPT-5 System Card](https://mlsanity.com/paper/2601.03267/) (2025)
- [Eureka: Human-Level Reward Design via Coding Large Language Models](https://mlsanity.com/paper/2310.12931/) (2023)
- [AirSim: High-Fidelity Visual and Physical Simulation for Autonomous Vehicles](https://mlsanity.com/paper/1705.05065/) (2017)
- [MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge](https://mlsanity.com/paper/2206.08853/) (2022)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [A Closer Look at Invalid Action Masking in Policy Gradient Algorithms](https://mlsanity.com/paper/2006.14171/) (2020)
- [Vision-Language Models as Success Detectors](https://mlsanity.com/paper/2303.07280/) (2023)
- [RoboCLIP: One Demonstration is Enough to Learn Robot Policies](https://mlsanity.com/paper/2310.07899/) (2023)
- [Behavior Trees in Robotics and AI: An Introduction](https://mlsanity.com/paper/1709.00084/) (2017)
- [Code as Reward: Empowering Reinforcement Learning with VLMs](https://mlsanity.com/paper/2402.04764/) (2024)
- [Neurosymbolic Reinforcement Learning and Planning: A Survey](https://mlsanity.com/paper/2309.01038/) (2023)
- [Latent Reward: LLM-Empowered Credit Assignment in Episodic Reinforcement Learning](https://mlsanity.com/paper/2412.11120/) (2024)
- [Learning to Follow Instructions in Text-Based Games](https://mlsanity.com/paper/2211.04591/) (2022)
- [Using Large Language Models to Automate and Expedite Reinforcement Learning with Reward Machine](https://mlsanity.com/paper/2402.07069/) (2024)
- [On the programming effort required to generate Behavior Trees and Finite State Machines for robotic applications](https://mlsanity.com/paper/2209.07392/) (2022)
- [Comparison between Behavior Trees and Finite State Machines](https://mlsanity.com/paper/2405.16137/) (2024)
- [Hierarchies of Reward Machines](https://mlsanity.com/paper/2205.15752/) (2022)
- [Designing Robust Cyber-Defense Agents with Evolving Behavior Trees](https://mlsanity.com/paper/2410.16383/) (2024)
- [Integrating Explanations in Learning LTL Specifications from Demonstrations](https://mlsanity.com/paper/2404.02872/) (2024)

## Abstract (arXiv)

Decomposing complex tasks into a sequence of simpler subtasks can improve learning efficiency for an autonomous agent. Reinforcement learning (RL) can be used to optimize agent policies to complete subtasks, but requires well-defined subtask rewards and benefits from action masking. Recent work uses large language models (LLMs) to automate reward shaping and action masking, however none of them fully address reactivity to subtask failure and modularity to varying objects for compositional tasks. To overcome these challenges, we develop masking reward behavior tree (MRBT), a symbolic structure used as a reactive and modular reward and action mask function. We design an MRBT template and derive logical specifications to construct and verify MRBTs for a sequence of object-interaction subtasks. Further, we develop an automated pipeline that uses an LLM to generate MRBTs robust to varying task objects, an SMT-solver to verify correctness of specifications, and a neurosymbolic RL loop to train agents on compositional tasks. Experiments demonstrate successful generation and refinement of five MRBTs, consistently improving training efficiency and task success rates over baselines and MRBTs without action masking. We further highlight three advantages of MRBTs: transferability, modularity, and verifiability.