💻 computer science

Active Reward Machine Inference From Raw State Trajectories

This paper proposes a method for learning reward machines directly from raw state and policy trajectories without access to rewards or labels, introducing an active learning framework to incrementally query trajectory extensions for improved data and computational efficiency.

Original authors: Mohamad Louai Shehab, Antoine Aspeel, Necmiye Ozay

Published 2026-04-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mohamad Louai Shehab, Antoine Aspeel, Necmiye Ozay

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to do a complex job, like cleaning a messy house. You don't just tell it "clean the house." You have to break it down: "Pick up the toys, then vacuum the rug, then wipe the table."

In the world of robotics, this "recipe" is often called a Reward Machine. It's like a flowchart that says, "If you see a toy, go to the 'Pick Up' stage. If you see a clean rug, go to the 'Vacuum' stage."

The Problem: The "Black Box" Mystery

Usually, a human expert has to draw this flowchart by hand. They have to decide exactly what counts as a "toy" (a label) and what the next step is. But in a real, messy world, this is incredibly hard. What if the robot sees a red ball? Is that a toy? Is it a decoration? If the human gets the flowchart wrong, the robot might get confused or do the wrong thing.

The big question this paper asks: Can we teach the robot to figure out the flowchart itself, just by watching it work, without us telling it what the steps are or what the "toys" look like?

The Solution: The Detective and the "What-If" Game

The authors propose a clever method to act like a detective. Here is how they do it, using simple analogies:

1. The "Memory" of the Robot

Imagine the robot is walking through a maze. It doesn't just see the floor tile it's standing on; it remembers the path it took to get there.

The Old Way: Researchers usually needed a "cheat sheet" (labels) telling them, "This tile is a wall," "That tile is a door."
The New Way: This paper says, "No cheat sheets needed!" We only have the robot's path (the raw data of where it went). We have to guess the rules of the maze just by watching where it goes.

2. The "Negative Example" (The "Wait, That's Wrong!" Moment)

The core of their method relies on finding contradictions.
Imagine you are guessing a secret code. You try a path, and the robot behaves one way. Then you try a slightly different path, and the robot behaves differently.

The Insight: If the robot acts differently, it must be in a different "mental state" or "stage" of the task.
The Analogy: Think of a video game. If you press "Jump" on a flat floor, you jump. If you press "Jump" while standing on a ladder, you might climb. The robot's reaction tells you that the "Floor" and the "Ladder" are different things in the game's logic, even if they look similar to a camera.

The paper uses a mathematical "puzzle solver" (called SAT) to find a flowchart that explains why the robot acted differently in those two situations.

3. The "Active Learning" Trick (The Smart Quiz)

Here is the biggest breakthrough. If you try to watch the robot walk every single possible path in a maze, you will run out of memory and time. The number of paths grows exponentially (like a tree branching out forever).

The authors realized you don't need to see every path. You just need to see the right paths.

The Analogy: Imagine you are trying to guess a friend's favorite movie.
- The Exhaustive Way: You ask them, "Do you like every single movie ever made?" (This takes forever and is impossible).
- The Active Way: You ask, "Do you like Action movies?" If they say yes, you immediately know to stop asking about Horror movies. You ask the specific question that splits the possibilities in half.

The paper's algorithm does exactly this. It looks at the thousands of possible paths the robot could take, picks the two that are most likely to confuse the current guesses, and asks, "Hey, if the robot took Path A vs. Path B, would it act differently?"

If the answer is Yes, it's a huge clue! It eliminates half of the wrong theories instantly.
If the answer is No, it's still useful, but less dramatic.

Why This Matters

This is a massive step forward because:

No Human Bias: We don't need to guess what the robot should be looking for. The robot figures out its own "vocabulary" (what counts as a "pickup" or a "drop-off").
Efficiency: By only asking the "smart questions" (active learning), the computer doesn't crash from trying to process too much data. It saves massive amounts of memory and time.
Real-World Ready: It moves us closer to robots that can learn complex, multi-step jobs just by watching an expert do them once, without needing a manual written in human language.

In a Nutshell

The paper is about teaching a robot to write its own instruction manual. Instead of us giving it the manual, we watch it work, spot the moments where it changes its behavior, and use a smart "quiz" to figure out the hidden rules of the game. It's like reverse-engineering a secret recipe just by tasting the final dish and asking, "What happened if I added salt instead of sugar?"

1. Problem Statement

The paper addresses the challenge of learning Reward Machines (RMs) directly from raw state trajectories in an information-scarce regime.

Context: Reward Machines are finite-state automata used to represent multi-stage tasks in robotics. They encode task progress via transitions triggered by high-level events (atomic propositions). Typically, RMs require a human expert to manually define the automaton structure and a labeling function that maps low-level states to high-level propositions.
The Gap: Existing methods for learning RMs usually assume access to ground-truth labels, reward signals, or an oracle. Some methods learn from demonstrations but still rely on a predefined labeling function.
The Core Problem: The authors aim to learn both the Reward Machine structure (transitions) and the Labeling Function simultaneously, using only:
1. The MDP model (states and transition dynamics).
2. A History Policy ( $\pi_h$ ), which captures the agent's behavior based on the sequence of observed states (without access to rewards, labels, or internal RM nodes).
Challenge: This problem is ill-posed because multiple RM/Labeling function pairs can induce the same optimal policy. The goal is to find a minimal RM that is policy-equivalent to the ground truth.

2. Methodology

The proposed framework consists of three main components: a formal problem definition, a SAT-based inference engine, and an active learning strategy to handle scalability.

A. Formalization: History Policy and Equivalence

History Policy ( $\pi_h$ ): Defined as the policy conditioned on the current state $s$ and the history of states $\tau$ . It represents the agent's behavior without explicit knowledge of the RM state.
Policy Equivalence: Two labeled RMs are considered equivalent if they induce the same history policy. The objective is to find the minimal RM within this equivalence class.
Negative Examples: The core constraint for learning comes from "negative examples." If two state trajectories $\tau$ and $\tau'$ lead to different actions in the history policy for the same state $s$ (i.e., $\pi_h(a|s, \tau) \neq \pi_h(a|s, \tau')$ ), then the underlying RM must have transitioned to different nodes after processing these trajectories.

B. SAT-Based Inference (Passive Learning)

The authors formulate the learning of the transition function ( $\delta_u$ ) and labeling function ( $L$ ) as a Boolean Satisfiability (SAT) problem.

Encoding: The transition function and labeling function are encoded as binary variables.
Constraints:
1. Functionality: Ensures $\delta_u$ and $L$ are valid functions (single output per input).
2. Anchoring: Fixes the initial node and the label of the initial state to break symmetry (permutation invariance).
3. Compatibility: Enforces that for every negative example pair $\{\tau, \tau'\}$ , the resulting RM nodes must be distinct ( $\delta^*_u(u_I, L(\tau)) \neq \delta^*_u(u_I, L(\tau'))$ ).
Sufficient Depth ( $l^*$ ): The paper proves Proposition 1, establishing that a sufficient depth $l^* = |S| \cdot u_{max}^2$ exists. If the SAT problem is satisfiable at this depth, it is satisfiable for all greater depths. This guarantees that the history policy contains enough information to uniquely identify the minimal RM structure (up to renaming).

C. Active Extension Algorithm

A major bottleneck is the exponential growth of state trajectories, making the storage of all negative examples infeasible (memory bottleneck). The authors propose an Active Learning approach to mitigate this:

Strategy: Instead of enumerating all trajectories, the algorithm iteratively queries specific trajectory extensions.
Volume Removal: The algorithm selects pairs of trajectories $\{\tau, \tau'\}$ ${τ, τ^{'}}$ that are most likely to split the current set of feasible solutions (candidate RMs) in half.
- It samples a subset of current feasible SAT solutions.
- It searches for trajectory pairs where half the candidates predict the same node (collapse) and the other half predict different nodes (separation).
Query: The system queries the history policy for these specific pairs. If they are confirmed as negative examples, they are added to the SAT constraints, drastically shrinking the hypothesis space.
Efficiency: This avoids storing millions of trajectories, focusing only on the "informative" ones needed to eliminate incorrect hypotheses.

3. Key Contributions

Information-Scarce Inference: First framework to learn both the RM structure and the labeling function from scratch using only raw state trajectories and a history policy, without rewards, labels, or node observations.
Theoretical Guarantees: Proved the existence of a sufficient depth ( $l^*$ ) for learning, ensuring that the problem is solvable given enough trajectory history.
Active Learning for RM: Introduced an active extension algorithm that uses a "volume removal" strategy to select informative trajectory pairs, significantly reducing memory and computational requirements compared to exhaustive enumeration.
Joint Learning: Formulated a unified SAT problem that solves for the automaton transitions and the state-labeling function simultaneously.

4. Experimental Results

The framework was tested on two grid-world tasks: Pick-and-Drop and Patrol (A→B→C→D).

Recovery Accuracy:
- In the Pick-and-Drop task, the algorithm recovered the ground-truth RM and labeling function (up to renaming) using a depth-9 history policy.
- In the Patrol task, it successfully recovered the correct structure and clustering of grid cells into propositions.
Active Learning Performance:
- Convergence: The active method converged to the ground-truth solution set in 100% of trials (for $N_{active}=100$ ) by depth 12, whereas a random baseline failed to find restrictive constraints even at depth 20.
- Memory Efficiency: The exhaustive approach required 24.76 GB to store ~414 million negative examples at depth 9. The active method reduced this to 0.147 GB (a ~168x reduction) by storing only ~0.29 million examples at a deeper horizon (depth 13).
- Runtime: The active method achieved a 2x speedup in total runtime (3544s vs 7185s) compared to the exhaustive baseline, primarily by reducing the time spent solving large SAT instances.

5. Significance and Future Work

Scalability: This work demonstrates that learning complex multi-stage task structures is feasible without human-provided symbolic knowledge or reward signals, provided the agent's policy is observable.
Memory in RL: The paper contributes to the theoretical understanding of memory in sequential decision-making, distinguishing between memory needed for partial observability and memory needed for task structure (temporal logic).
Limitations & Future Directions:
- Current formulation assumes discrete state spaces and access to the optimal history policy.
- Future work aims to handle estimation errors in the history policy (from finite data) and extend the framework to continuous/perceptual domains (e.g., learning labels from images).
- Investigating termination criteria based on policy equivalence rather than just structural minimality.

In summary, this paper provides a rigorous, scalable, and automated method for reverse-engineering the logical structure of complex robotic tasks directly from observed behavior, removing a significant barrier to the deployment of reward-based learning in real-world environments.