Stratifying Reinforcement Learning with Signal Temporal Logic

1. Problem Statement

Modern machine learning embeddings, particularly those from Deep Reinforcement Learning (DRL) agents, often violate the "manifold hypothesis," which assumes data lies on a single smooth manifold of fixed dimension. Instead, these spaces are stratified, consisting of manifolds of varying dimensions glued together.

The paper addresses two main gaps:

1. Theoretical Gap: There is a lack of understanding regarding how the stratified structure of an ambient data domain (e.g., a game environment) induces a corresponding stratification in the latent embedding space of a neural network.
2. Methodological Gap: Existing tools for analyzing embedding geometry often fail to capture the complex, multi-dimensional structures inherent in decision-making tasks constrained by temporal logic.

The authors aim to bridge control theory, machine learning, and algebraic topology by modeling DRL decision spaces as poset-stratified spaces and using Signal Temporal Logic (STL) robustness as a reward function to shape these spaces.

2. Methodology

The paper proposes a framework combining Poset Stratification Theory, Signal Temporal Logic (STL), and Numerical Topological Data Analysis (TDA).

A. Theoretical Framework: Poset Stratification

• Definition: A stratified space is defined via a continuous map $q: X \to P$, where $X$ is the topological space (state or embedding space) and $P$ is a partially ordered set (poset). The pre-image of each element in $P$ is a manifold.
• Application to Games:
  • Minigrid Games: The authors model 2D grid environments as cell complexes. They demonstrate that policies (routes to goals) induce a tree-stratified space, where the tree represents the hierarchy of states leading to the goal.
  • Coin Collection (Space-Time): For games involving time and velocity constraints, the state space is modeled as a "light cone" poset. The authors introduce an overlap poset to handle intersections of future light cones, proving that trajectory spaces can be stratified by the order complex of these posets.

B. STL as a Reward and Stratification Tool

• STL Semantics: The authors reinterpret STL atomic predicates as membership tests within stratified spaces.
• Rug Functions & Tubular Neighborhoods: They utilize "rug functions" (semi-algebraic functions) to define tubular neighborhoods around strata.
• Reward Mechanism: Instead of sparse binary rewards, the DRL agent is trained using the robustness score of an STL formula as the reward. This forces the agent to not just satisfy the logic (e.g., "eventually reach the green square") but to do so with a margin of safety, effectively shaping the geometry of the decision space.

C. Numerical Detection of Stratification

To analyze high-dimensional latent spaces, the authors propose and compare several dimensionality estimation techniques:

• Volume Growth Transform (VGT): Analyzes the log-log growth of volume within a ball of radius $r$ around a point. The slope indicates local dimension.
• VGT-dot: A novel contribution defined as the derivative of the smoothed VGT curve. This feature is proposed to reveal local-to-global stratification signatures (e.g., detecting where a 2D room transitions to a 1D corridor).
• HADES & DIC: The authors compare their methods against HADES (a singularity detector) and Dimension Induced Clustering (DIC).

3. Key Contributions

1. Novel Correspondence Principle: The paper establishes a theoretical link showing that most STL formulas induce a stratification of space-time. It argues that the reward landscape defined by STL robustness naturally creates a stratified structure in the agent's latent space.
2. VGT-dot Signature: The introduction of the VGT-dot derivative as a computationally efficient feature for detecting stratification structures in high-dimensional data, capable of identifying "pinch points" and singularities where local dimension changes.
3. Poset-Stratified Game Modeling: A formalization of Minigrid and space-time games as poset-stratified spaces, providing a rigorous topological language for describing RL policies and state transitions.
4. Empirical Validation: Application of these theories to a Transformer-based DRL agent, demonstrating that the latent embeddings of the agent exhibit clear stratified structures (specifically "hourglass" shapes) corresponding to the temporal logic constraints of the task.

4. Experimental Results

The authors trained a Transformer-XL agent on Minigrid environments with two specific STL reward functions:

1. "Eventually in the green square": The agent must reach a target within a specific time window.
2. "Eventually in green AND always not in red": A more complex constraint avoiding obstacles.

Findings:

• Latent Space Geometry: UMAP projections of the 256-dimensional token embeddings revealed distinct point clouds.
  • The "Eventually" task produced an hourglass-shaped structure. The "neck" of the hourglass corresponded to the transition period when the temporal operator became active (the square filled).
  • The "Always not" task mirrored this structure, showing two lobes connected by a narrow region.
• Stratification Detection:
  • VGT-dot: Successfully identified the "neck" of the hourglass as a region of changing dimensionality, distinguishing it from the bulk of the 2D-like room and 1D-like corridor.
  • HADES: When applied to a dimension-reduced (100D) version of the data, HADES flagged the "neck" points as non-manifold (singularities) with statistical certainty, confirming the stratified nature of the space.
  • Clustering: Agglomerative clustering based on VGT-dot features successfully separated the space into core (low dimension/singular) and boundary (high dimension) regions.

5. Significance and Conclusion

• Theoretical Impact: The work provides a rigorous mathematical foundation for understanding why DRL embeddings are not simple manifolds but complex stratified spaces. It bridges the gap between the discrete logic of STL and the continuous geometry of neural networks.
• Practical Utility: The proposed VGT-dot technique offers a scalable, computationally efficient method for analyzing high-dimensional embeddings without the memory constraints of methods like HADES.
• Future Directions: The paper suggests that understanding these stratifications can lead to better interpretability of DRL agents, improved robustness verification, and the development of new algorithms that explicitly account for the geometry of decision spaces.

In summary, the paper demonstrates that by framing RL tasks through the lens of Signal Temporal Logic and stratification theory, one can uncover the hidden geometric structure of neural network embeddings, revealing how agents navigate complex, multi-dimensional decision landscapes.