TLDR: A new research paper introduces “Symmetry-Aware Transformer Training” to improve how AI models solve complex planning problems. Traditional transformer models like PlanGPT struggle with the arbitrary naming of objects and the order of information in planning tasks, limiting their ability to generalize to new, harder problems. The new approach uses a novel contrastive learning objective and architectural adjustments to make transformers ‘symmetry-aware’, enabling them to learn more efficiently and extrapolate better to unseen problem sizes. This significantly boosts performance and training stability in automated planning.
Automated planning is a crucial area in artificial intelligence, focusing on how intelligent agents can devise sequences of actions to achieve specific goals. Think of it like a robot figuring out the steps to assemble a product or navigate a complex environment. Recently, transformer models, which have revolutionized fields like natural language processing, have shown promise in this domain. However, their application in automated planning has faced significant hurdles, particularly when problems become more complex or involve unfamiliar scenarios.
A notable effort in this space is PlanGPT, a state-of-the-art transformer model designed to predict plans from formal problem descriptions. While innovative, PlanGPT encounters several key limitations that hinder its ability to generalize from simpler training examples to more challenging, unseen problems. These limitations primarily stem from the inherent ‘symmetries’ within planning tasks.
Challenges with Existing Transformer Models in Planning
One major issue is how planning problems are represented. Objects in a planning task, like ‘truck5’ or ‘loc3’, are essentially just identifiers; their names carry no intrinsic meaning beyond distinguishing them. This leads to a combinatorial explosion of equivalent representations for the same problem. For instance, if you have eight objects, there are over 40,000 ways to assign them names from a vocabulary, all describing the identical planning task. PlanGPT struggles to efficiently learn that these are all the same problem, leading to poor learning efficiency.
Furthermore, PlanGPT can inadvertently ‘memorize’ information from object names. In some domains, like ‘Visitall’ where locations are named ‘loc-x1-y2’ (indicating coordinates), PlanGPT might rely on these names rather than the underlying problem structure. This prevents it from generalizing to new locations not seen during training.
Another challenge arises from the arbitrary order of ‘atoms’ (basic propositions about the world, like ‘at(package1, location2)’) in the initial and goal states. The order doesn’t change the problem, but PlanGPT, being a decoder-only model, uses learned positional encodings that can overfit to specific atom orders. This makes it difficult for the model to generalize to larger problems where atom orders might differ.
Finally, the reliance on learned positional encodings also limits generalization to longer plans. If a model is trained on short plans, it may not have learned encodings for positions that appear in much longer, more complex plans, thus failing to solve them.
Introducing Symmetry-Aware Transformer Training
To overcome these limitations, researchers have proposed a novel approach called “Symmetry-Aware Transformer Training.” This method introduces a new way of training transformers that explicitly accounts for the symmetries in planning problems, guiding the model to learn representations that are robust to variations in object names and atom orders. You can read the full research paper here: Symmetry-Aware Transformer Training for Automated Planning.
The core of this new training objective involves using pairs of symmetric planning problems. For any given problem, two versions are generated that are identical in structure but have different, randomly assigned object names. The model is then trained to behave equivalently for both versions, essentially learning to ignore the arbitrary naming conventions.
This is achieved through two complementary ‘contrastive loss’ components:
- Attention Loss: This component ensures that the model’s internal ‘attention scores’ (which determine how different parts of the input relate to each other) are identical for corresponding parts of the symmetric input pairs. This means the model learns the same underlying ‘algorithm’ regardless of the object names.
- Hidden State Loss: This component encourages the learned representations (hidden state vectors) for corresponding tokens to be similar, ensuring that the model captures the problem’s structure independently of specific object names.
Beyond the training objective, architectural improvements were also implemented. For processing the initial and goal states, an encoder-only or encoder-decoder architecture is used without any positional encodings. This makes the encoder inherently robust to the arbitrary order of input atoms. Additionally, a compositional tokenization scheme is employed, where each atom (like ‘at(truck5, loc3)’) is encoded as a single embedding, rather than splitting it into multiple tokens, further enhancing permutation-equivariance.
For plan generation, where action order is crucial, the decoder explicitly omits positional encoding (a technique known as NoPE). This has been shown to improve generalization to unseen sequence lengths, addressing the limitation of learned positional encodings.
Also Read:
- Enhancing World Models for Robotics Through Novelty Detection
- Boosting Efficiency in Learning Logic Programs
Promising Results and Future Outlook
The new symmetry-aware training approach was evaluated across several planning domains, including Blocksworld, Gripper, Visitall, and Logistics. The results demonstrate significant improvements in the model’s ability to extrapolate to harder planning problems compared to PlanGPT. While PlanGPT often failed entirely on larger instances, the new models were able to solve them to varying extents, showcasing better generalization capabilities.
The contrastive loss also proved beneficial for training stability, with models trained using this objective exhibiting lower validation losses and fewer instances of training divergence compared to those without it. This indicates a more robust and reliable learning process.
Despite these advancements, challenges remain. The models still face limitations in consistently generating valid plans for problems significantly larger than those encountered during training, especially in complex domains like Logistics. The fixed vocabulary size also poses a practical constraint for very large tasks. Future research will focus on identifying these architectural or representational bottlenecks and exploring techniques that are independent of fixed token vocabularies to tackle even larger and more complex planning problems.
