🤖 machine learning

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

This paper introduces the Masking Reward Behavior Tree (MRBT), a neurosymbolic framework that leverages large language models and SMT solvers to automatically generate verifiable, modular, and reactive reward shaping and action masking functions, thereby significantly improving reinforcement learning efficiency and success rates for compositional tasks involving varying objects.

Original authors: Nicholas Potteiger, Ankita Samaddar, Taylor T. Johnson, Xenofon Koutsoukos

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nicholas Potteiger, Ankita Samaddar, Taylor T. Johnson, Xenofon Koutsoukos

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to clean a messy house. If you just tell the robot, "Clean the house," and give it a big reward only when the whole job is done, the robot will likely get confused. It might sweep the floor, then drop the vacuum, then try to wash the windows, and never actually finish. It's like trying to learn a complex dance by only getting a prize at the very end of the song; you won't know which steps were right or wrong along the way.

This paper introduces a smarter way to teach robots (or autonomous agents) how to handle complex, multi-step jobs. The authors call their solution MRBT (Masking Reward Behavior Tree). Here is how it works, broken down into simple concepts:

1. The Problem: The "Black Box" of Rewards

In traditional robot training (Reinforcement Learning), you have to manually design a "reward system." You have to decide exactly what the robot gets a "point" for and what it gets a "penalty" for.

The Hard Part: If the robot fails a small step (like dropping a key), a simple reward system might not know to tell the robot to go back and try again. It might just keep wandering aimlessly.
The Modularity Issue: If you change the task slightly (e.g., instead of a red key, use a blue key), you often have to rewrite the whole reward system from scratch.

2. The Solution: A "Smart Coach" with a Checklist

The authors use a Behavior Tree. Think of this as a flowchart or a checklist that the robot follows. It breaks the big job into small, manageable steps (subtasks).

The "Masking" Trick: Imagine the robot has a hand full of tools. Sometimes, it needs to hide (mask) the tools it doesn't need right now. For example, if the robot is supposed to walk to a door, the system "masks" (hides) the "pick up" button so the robot doesn't accidentally try to pick up the air. This stops the robot from wasting time on useless actions.
The "Reward" Trick: The system gives small rewards for finishing each step on the checklist, not just the final goal. If the robot drops the key, the system immediately says, "Oh no, you dropped it! Go back to the start of this step," rather than letting the robot wander off.

3. The Magic Ingredient: The AI "Architect" (LLM)

Designing these checklists and rules by hand is hard. So, the authors used a Large Language Model (LLM)—the same kind of AI that powers chatbots—to do the heavy lifting.

You give the AI a template (like a blank form) and a description of the task (e.g., "Get the key from the red room").
The AI fills in the blanks: It writes the logic for "Is the door open?", "Is the key near?", and decides which buttons to hide at each step.
The Safety Net (SMT Solver): Since AI can make mistakes, the authors added a "logic checker" (an SMT solver). This is like a strict math teacher who checks the AI's work. If the AI writes a rule that doesn't make sense (e.g., "The door is open" but the robot is still far away), the checker catches it, tells the AI, "This is wrong," and asks it to try again.

4. The Results: A Robot That Learns Faster

The team tested this on two different "playgrounds":

MiniGrid: A simple grid world where a robot has to find keys, open doors, and reach a goal.
MuJoCo Fetch: A more realistic simulation of a robotic arm picking up blocks.

They found that robots trained with their MRBT system:

Learned much faster: They reached the goal in fewer tries.
Were more successful: In the hardest tasks, they succeeded over 80% of the time, while other methods struggled below 70%.
Handled mistakes better: If the robot dropped an item, the system immediately guided it back to fix the error, rather than letting it get lost.

5. Why This Matters (According to the Paper)

The authors highlight three main superpowers of their system:

Transferability: They trained a robot in a simple grid world and it worked surprisingly well when moved to a realistic drone simulator (AirSim). It's like learning to drive in a parking lot and then being able to drive on a real highway.
Modularity: If you want to add a new step to the task (like "wash the windows" after "clean the floor"), you can just plug it into the tree without rebuilding the whole system. It's like adding a new chapter to a book without rewriting the whole story.
Verifiability: Because they used the "logic checker," they can mathematically prove that the rules the AI generated are correct, rather than just hoping they work.

In short: The paper shows how to use AI to automatically write a "smart coach" for robots. This coach breaks big jobs into small steps, hides useless buttons, and instantly corrects mistakes, making robots learn complex tasks much faster and more reliably than before.

Technical Summary: Reward Shaping and Action Masking for Compositional Tasks using Behavior Trees and LLMs

1. Problem Statement

Autonomous agents often face complex tasks that require decomposing a high-level goal into a sequence of simpler subtasks, such as navigating to an object, manipulating it, and proceeding to a goal. While Reinforcement Learning (RL) can optimize policies for these subtasks, it suffers from low sample efficiency without well-defined reward functions and action masks. Manually designing these components is challenging because it requires:

Identifying appropriate subtasks.
Ensuring reactivity to subtask failures (e.g., if an agent drops a key, it must backtrack to pick it up again rather than proceeding).
Maintaining modularity across varying task objects (e.g., the same logic must apply whether the object is a "red key" or a "blue key").

Existing automated approaches using Large Language Models (LLMs) or Vision-Language Models (VLMs) often generate reward code or finite-state machines but fail to fully address reactivity to failures, modularity for compositional tasks, or the joint optimization of reward and action masking.

2. Methodology

The authors propose Masking Reward Behavior Trees (MRBTs), a symbolic structure that serves as both a reward shaping function and an action masking function. The methodology integrates three core components: a behavior tree template, logical specifications for verification, and an automated pipeline using an LLM and an SMT solver.

2.1. Masking Reward Behavior Tree (MRBT)

An MRBT is a Behavior Tree (BT) where leaf nodes are Masking Behavior Reward Machines (MBRMs).

Structure: The BT coordinates execution via control nodes (Sequence → and Fallback ?).
Leaves: Each leaf is an MBRM that outputs a scalar reward and a discrete action mask based on the agent's state and task.
Reactivity: The BT structure enables backtracking. If a subtask fails (e.g., an interaction fails), the BT reverts to a previous state, ensuring the agent retries the necessary step rather than receiving a reward for an incomplete sequence.
Modularity: The MRBT uses placeholders for subtasks and logic formulas, allowing it to adapt to different object attributes (e.g., colors) within a defined task space.

2.2. MRBT Template

The authors design a specific template for sequential object-interaction tasks. For a sequence of $k$ subtasks, the template includes:

Logic Formulas: For each subtask $i$ , it defines a completion formula ( $\psi_i$ ) and an object proximity formula ( $\phi_i$ ).
Action Masks: Two masks per subtask: $\eta_{i,1}$ for navigation and $\eta_{i,2}$ for interaction.
Execution Logic: The BT executes a sequence of MBRMs. If a subtask is complete, it moves to the next; if the agent is near the object, it enables interaction; otherwise, it enables navigation. If a subtask fails (state transitions to Failure), the BT backtracks.

2.3. Automated Generation and Verification Pipeline

To automate the creation of MRBTs, the authors developed a pipeline (Figure 1) involving:

LLM Generation: An LLM (ChatGPT-5) takes a task description, the MRBT template, and environment predicates as input. It generates the specific subtasks, logic formulas ( $\psi_i, \phi_i$ ), and action masks ( $\eta_{i,1}, \eta_{i,2}$ ).
SMT Verification: A symbolic model of the environment dynamics ( $E_{sym}$ $E_{sy m}$ ) is constructed in an SMT solver (Z3). The solver checks the generated logic formulas against three specifications:
- Completion Correctness: Ensures that if the final task is complete, all subtasks were completed.
- Object Proximity: Ensures the agent is near the object immediately before subtask completion.
- Non-regressive Maximal Reward: Ensures that once a subtask is completed, it remains completed (no regression) to achieve the maximal reward.
Refinement: If the SMT solver finds a counterexample (e.g., a trajectory where the agent completes the task without satisfying a subtask condition), it generates a debug prompt. The LLM is re-prompted with this counterexample to refine the logic formulas until they are satisfiable.
Neurosymbolic RL Training: The verified MRBT is integrated into a neurosymbolic RL loop (using PPO) to train the agent. The MRBT provides the reward signal and restricts the action space at each step.

3. Key Contributions

MRBT Framework: Introduction of a symbolic structure that jointly outputs rewards and action masks, specifically designed for reactivity and modularity in compositional tasks.
MRBT Template & Specifications: A systematic template for constructing MRBTs for sequential object-interaction tasks, accompanied by logical specifications to verify correctness (completion, proximity, and non-regression) using an SMT solver.
Automated Pipeline: A system that uses an LLM to generate MRBTs and an SMT solver to verify and refine them, ensuring robustness across a task space with varying objects.
Comprehensive Evaluation: Experiments across five task spaces in two environments (MiniGrid and MuJoCo Fetch) under both deterministic and stochastic dynamics.
Advantages Analysis: Demonstration of MRBTs' transferability (to AirSim), modularity (scaling better than Hierarchical Reward Machines), and verifiability (theoretical guarantees via Z3).

4. Experimental Results

The authors evaluated the MRBT approach against two baselines:

Task: A binary reward given only upon total task completion.
Procedure: A reward machine based on VLM-CaR that awards rewards for subtask completion but lacks reactive backtracking and action masking.
RBT (Ablation): An MRBT without action masking.

Key Findings:

Training Efficiency & Success: MRBTs consistently achieved higher training efficiency and task success rates compared to baselines and the RBT ablation.
Complexity Handling: In the most complex task space (LockedRoom, 4 subtasks, 36 variants), MRBTs achieved an average success rate of $\ge 80\%$ , whereas baselines achieved $\le 70\%$ .
Stochasticity: MRBTs maintained high performance in stochastic environments (e.g., keys dropping randomly), demonstrating the value of reactive backtracking.
Transferability: Policies trained with MRBTs in MiniGrid successfully transferred to a realistic quadcopter simulator (AirSim), achieving a 97% success rate in deterministic settings and 94% in stochastic settings, outperforming all ablations.
Modularity: The MRBT structure scales more modularly than Hierarchical Reward Machines (HRMs); adding a subtask to an MRBT requires $O(1)$ storage, whereas HRMs require $O(k)$ storage for back edges.
Verification: The SMT solver successfully refined generated logic formulas (e.g., correcting proximity thresholds or persistence conditions) in under 10 minutes for most tasks.

5. Significance and Claims

The paper claims that MRBTs offer a robust solution for automating reward shaping and action masking in compositional tasks. The significance lies in three main advantages:

Transferability: The symbolic nature of MRBTs allows policies trained in simulation to transfer effectively to realistic simulators (AirSim) and varying task objects.
Modularity: Unlike finite-state machines or reward machines, the tree structure of MRBTs allows for easier scaling and modification as task complexity increases.
Verifiability: By integrating SMT solvers, the approach provides theoretical guarantees that the reward and action mask logic are correct with respect to the task specifications, a feature not present in purely data-driven or demonstration-based approaches.

The authors conclude that while the approach requires an SMT solver for verification, the resulting MRBTs significantly improve learning efficiency and task success, particularly in complex, multi-step environments where reactivity to failure is critical. Future work is noted to explore generalizing the MRBT template further.