STLCG

JAX toolbox for differentiable Signal Temporal Logic robustness

License

BSD-2-Clause

Overview

STLCG (distributed as the JAX library stljax) is a toolbox for computing the robustness of Signal Temporal Logic (STL) formulas using computation graphs. STL is a formal language for expressing spatio-temporal requirements such as "enter the target region within 10 seconds and stay for 5 seconds", and its robustness measure quantifies how strongly a signal satisfies or violates a specification. stljax reimplements the original 2019 Stanford ASL PyTorch stlcg in JAX, replacing the sequential recurrent evaluation with a fully parallelizable masking-based formulation that runs up to ~1000x faster on GPU. Because the whole pipeline is differentiable, STL specifications can be embedded directly into gradient-based trajectory optimization, reinforcement-learning reward shaping, and parameter mining. Published at IEEE Robotics and Automation Letters (RA-L) 2025; pip-installable as stljax.

stljax lets robotics engineers turn human-readable temporal-logic mission specifications into differentiable objective and constraint terms. You build STL formulas from Expressions and Predicates, evaluate their robustness over a signal or trajectory, and—because everything is implemented in JAX—get batched evaluation via vmap and exact gradients for use inside optimizers and learning loops. The key design choice is a masking-based reformulation of the temporal operators (Always/Eventually/Until) that parallelizes what used to be a sequential recurrent computation, giving order-of-magnitude speedups while also distributing gradients correctly across all min/max operations.

Why it matters

Signal Temporal Logic gives robotics a precise, human-readable way to specify temporal and spatial mission requirements, but for years its robustness computation was sequential, slow, and awkward to differentiate—keeping STL out of modern learning and optimization pipelines. stljax removes that barrier: a parallelized, autodiff-native implementation makes STL fast enough for real-time control and large-batch training, so specifications can act directly as planner objectives, RL rewards, or verification signals.

Problem it solves

Prior STL robustness tools relied on sequential recurrent computation that was slow and hard to differentiate, which limited the use of formal temporal-logic specifications inside gradient-based trajectory optimization and learning.

Key takeaways

• Masking-based parallel STL robustness, up to ~1000x faster than recurrent methods on GPU
• JAX-native: vmap batching and exact gradients, with correct gradient distribution across min/max
• Expression and Predicate APIs for building Always/Eventually/Until formulas
• Differentiable specs usable as objectives/constraints in trajectory optimization and RL
• pip-installable (stljax); RA-L 2025; successor to Stanford ASL's PyTorch stlcg

Who it is for

• Motion-planning and trajectory-optimization researchers
• Control and formal-methods researchers
• Reinforcement-learning reward designers
• Safe-autonomy and verification engineers

Limitations

• Requires familiarity with Signal Temporal Logic to author useful specifications
• Smoothing/approximation of the robustness function introduces accuracy/gradient trade-offs
• JAX dependency and its toolchain
• A specification-evaluation library—does not include a planner or controller itself