Learning Top- \boldsymbolK Subtask Planning Tree Based on Discriminative Representation Pretraining for Decision-making

Decomposing complex real-world tasks into simpler subtasks and devising a subtask execution plan is critical for humans to achieve effective decision-making. However, replicating this process remains challenging for AI agents and naturally raises two questions: 1) How to extract discriminative knowledge representation from priors? 2) How to develop a rational plan to decompose complex problems? To address these issues, we introduce a groundbreaking framework that incorporates two main contributions. First, our multiple-encoder and individual-predictor regime goes beyond traditional architectures to extract nuanced task-specific dynamics from datasets, enriching the feature space for subtasks. Second, we innovate in planning by introducing a top-K subtask planning tree generated through an attention mechanism, which allows for dynamic adaptability and forward-looking decision-making. Our framework is empirically validated against challenging benchmarks BabyAI including multiple combinatorially rich synthetic tasks (e.g., GoToSeq, SynthSeq, BossLevel), where it not only outperforms competitive baselines but also demonstrates superior adaptability and effectiveness in complex task decomposition.