Unified system speeds up multi-task robot training with logic
Jaxolotl: A Unified High-Performance Benchmark Suite for LTL-Based Multi-Task RL
Machine LearningArtificial Intelligence
Summary
Teaching robots or agents to follow many different instructions is hard, especially when instructions involve time or order. The authors introduce Jaxolotl, a software suite that helps train and test these agents much faster and more fairly by using a common setup and tricks to speed up computations. Using this tool, they compared various existing methods and found that some can think ahead but slow down with more details, while others handle many tasks better but think less deeply.
What this means in practice
- •For robot developers: Train robots on complex, time-dependent tasks using a standard, fast toolkit that enables fair comparison of instruction-following methods.
- •For ai algorithm engineers: Evaluate and compare multi-task reinforcement learning algorithms more efficiently and at larger scale with a high-speed unified benchmark.
Authors
Mathias Jackermeier, Jacques Cloete, Alessandro Abate
Abstract
Training agents to follow arbitrary instructions is an important goal of multi-task reinforcement learning (RL). Linear temporal logic (LTL) provides a precise and structured formalism for specifying instructions to agents, and has been successfully adopted for training generalist multi-task policies. However, differences in implementations, task distributions, and evaluation protocols make existing methods difficult to compare, while high computational costs limit the scale and statistical reliability of experiments. We introduce Jaxolotl, a unified high-performance benchmark suite for multi-task LTL-RL to address these concerns. Jaxolotl provides a modular, end-to-end JAX implementation of six representative algorithms and four environments, together with newly curated task suites and a standardised, statistically robust evaluation protocol. By precompiling symbolic task representations into static arrays, Jaxolotl enables fully JIT-compiled training and evaluation, achieving end-to-end speedups of up to $220\times$ and supporting controlled comparisons at substantially greater experimental scale. We use this framework to systematically evaluate existing approaches, revealing complementary strengths and limitations: general methods capable of non-myopic reasoning struggle as the number of propositions grows, while methods with stronger scaling rely on environment-specific assumptions and suffer from myopia.