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

• URL: http://arxiv.org/abs/2605.05795v1
• Date: Thu, 07 May 2026 07:33:08 GMT
• Title: Reward Shaping and Action Masking for Compositional Tasks using Behavior Trees and LLMs
• Abstract summary: Reinforcement learning can be used to optimize agent policies to complete subtasks.<n>Recent work uses large language models (LLMs) to automate reward shaping and action masking.<n>We develop masking reward behavior tree (MRBT), a symbolic structure used as a reactive and modular reward and action mask function.
• Score: 3.285378336473673
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Learning to Adapt: Representation-Based Reinforcement Learning for Multi-Task Skill Transfer [5.787601808065202]
  RepMT-SAC is a framework for multi-task RL that enables efficient knowledge sharing and robust transfer to new tasks.<n>We evaluate RepMT-SAC on trajectory-following tasks across in-distribution and out-of-distribution contexts, demonstrating that it outperforms baselines by up to 30%.
  arXiv  Detail & Related papers  (2026-06-11T04:33:03Z)
• Decomposing the Basic Abilities of Large Language Models: Mitigating Cross-Task Interference in Multi-Task Instruct-Tuning [58.90602461025722]
  Multi-task instruct-tuning in large language models (LLMs) has been largely driven by multi-task instruct-tuning.<n>The training paradigm suffers from a key issue, named cross-task interference, due to conflicting gradients over shared parameters among different tasks.<n>We propose BADIT that decomposes LLM parameters into high-singular-value LoRA experts representing basic abilities.<n>We conduct extensive experiments on the SuperNI benchmark with 6 LLMs, and empirical results demonstrate that BADIT can outperform SOTA methods and mitigate the degree of cross-task interference.
  arXiv  Detail & Related papers  (2026-05-07T05:08:58Z)
• Efficient Knowledge Transfer in Multi-Task Learning through Task-Adaptive Low-Rank Representation [11.955971931186006]
  Pre-trained language models struggle with emerging tasks unseen during training in real-world applications.<n>We propose Task-Adaptive Low-Rank Representation (TA-LoRA), an MTL method built on prompt tuning.<n>Experiments on 16 tasks demonstrate that TA-LoRA achieves state-of-the-art performance in full-data and few-shot settings.
  arXiv  Detail & Related papers  (2025-04-20T06:33:19Z)
• Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models [45.12916211850169]
  behavior foundation models (BFMs) enable multi-modal, human-like control for humanoid agents.<n>"Task Tokens" are a method to effectively tailor BFMs to specific tasks while preserving their flexibility.<n>We show that Task Tokens offer a promising approach for adapting BFMs to specific control tasks while retaining their generalization capabilities.
  arXiv  Detail & Related papers  (2025-03-28T21:28:13Z)
• Layer-Aware Task Arithmetic: Disentangling Task-Specific and Instruction-Following Knowledge [12.367471198090655]
  Task Arithmetic (TA), which combines task vectors derived from fine-tuning, enables multi-task learning and task forgetting but struggles to isolate task-specific knowledge from general instruction-following behavior.<n>We propose Layer-Aware Task Arithmetic (LATA), a novel approach that assigns layer-specific weights to task vectors based on their alignment with instruction-following or task-specific components.
  arXiv  Detail & Related papers  (2025-02-27T15:22:14Z)
• Guiding Multi-agent Multi-task Reinforcement Learning by a Hierarchical Framework with Logical Reward Shaping [16.5526277899717]
  This study aims to design a multi-agent cooperative algorithm with logic reward shaping. Experiments have been conducted on various types of tasks in the Minecraft-like environment.
  arXiv  Detail & Related papers  (2024-11-02T09:03:23Z)
• Supervised Fine-Tuning Achieve Rapid Task Adaption Via Alternating Attention Head Activation Patterns [47.57912649802414]
  We study the process that the SFT process adapts LLMs to downstream tasks via the perspective of attention patterns. We find that LLMs selectively activate task-specific attention heads during SFT; (2) activation patterns for complex tasks are combinations of basic task patterns; and (3) changes in a few parameters can significantly impact activation patterns after SFT on a small number of samples.
  arXiv  Detail & Related papers  (2024-09-24T07:34:50Z)
• Re-TASK: Revisiting LLM Tasks from Capability, Skill, and Knowledge Perspectives [54.14429346914995]
  Chain-of-Thought (CoT) has become a pivotal method for solving complex problems with large language models (LLMs)<n>This paper introduces the Re-TASK framework, a novel theoretical model that revisits LLM tasks from capability, skill, and knowledge perspectives.<n> Experiments across diverse domains demonstrate the effectiveness of Re-TASK.
  arXiv  Detail & Related papers  (2024-08-13T13:58:23Z)
• TaskBench: Benchmarking Large Language Models for Task Automation [82.2932794189585]
  We introduce TaskBench, a framework to evaluate the capability of large language models (LLMs) in task automation. Specifically, task decomposition, tool selection, and parameter prediction are assessed. Our approach combines automated construction with rigorous human verification, ensuring high consistency with human evaluation.
  arXiv  Detail & Related papers  (2023-11-30T18:02:44Z)
• CL-MAE: Curriculum-Learned Masked Autoencoders [49.24994655813455]
  We propose a curriculum learning approach that updates the masking strategy to continually increase the complexity of the self-supervised reconstruction task. We train our Curriculum-Learned Masked Autoencoder (CL-MAE) on ImageNet and show that it exhibits superior representation learning capabilities compared to MAE.
  arXiv  Detail & Related papers  (2023-08-31T09:13:30Z)
• Semantically Aligned Task Decomposition in Multi-Agent Reinforcement Learning [56.26889258704261]
  We propose a novel "disentangled" decision-making method, Semantically Aligned task decomposition in MARL (SAMA) SAMA prompts pretrained language models with chain-of-thought that can suggest potential goals, provide suitable goal decomposition and subgoal allocation as well as self-reflection-based replanning. SAMA demonstrates considerable advantages in sample efficiency compared to state-of-the-art ASG methods.
  arXiv  Detail & Related papers  (2023-05-18T10:37:54Z)
• Plan, Eliminate, and Track -- Language Models are Good Teachers for Embodied Agents [99.17668730578586]
  Pre-trained large language models (LLMs) capture procedural knowledge about the world. Plan, Eliminate, and Track (PET) framework translates a task description into a list of high-level sub-tasks. PET framework leads to a significant 15% improvement over SOTA for generalization to human goal specifications.
  arXiv  Detail & Related papers  (2023-05-03T20:11:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.