Probabilistic Recurrent Intention Switching Model

Technical Summary: Probabilistic Recurrent Intention Switching Model (PRISM)

Problem Statement
Standard Inverse Reinforcement Learning (IRL) assumes a single, stationary reward function drives an agent's behavior. However, real-world agents—both biological (e.g., mice navigating labyrinths) and artificial (e.g., teleoperated robots)—frequently switch between distinct goals or strategies within a single episode. Existing multi-intention IRL approaches struggle to model these switches effectively:

• Dynamic IRL (DIRL) assumes smooth, continuous reward variation, conflicting with evidence of discrete strategy switching and requiring computationally expensive inference ($T$ Bellman solves).
• Hierarchical Inverse Q-Learning (HIQL) models intention transitions as a memoryless Markov chain, failing to capture history-dependent switches driven by factors like fatigue or satiation. Its inference scales poorly ($O(nK^2)$) due to the Baum–Welch algorithm.
• SWIRL incorporates history via state augmentation, but this causes the state space to explode exponentially ($|S|^L$), limiting practical history lengths.

Consequently, no existing method naturally scales to large, non-Markovian environments while recovering interpretable, discrete intentions.

Methodology: PRISM
The authors propose the Probabilistic Recurrent Intention Switching Model (PRISM), a framework that replaces fixed transition models and manual state augmentation with a lightweight recurrent neural network (RNN).

1. Model Architecture:

   • Hidden-Intention MDP (HI-MDP): The environment is modeled as an MDP with a latent intention space $Z = \{1, \dots, K\}$.
   • Recurrent Gating: A parametric function $f_\theta$ (an RNN, LSTM, or Transformer) maps the observation history $\phi_{1:t}$ to a soft probability distribution over intentions at each timestep. This allows the model to capture complex, history-dependent switching mechanisms without explicit state augmentation.
   • Per-Intention Rewards: Each intention $k$ is associated with a specific reward function $r_k$. The expert's policy under intention $k$ is assumed to be a Boltzmann-optimal policy derived from $r_k$.

2. Inference and Optimization:

   • EM Decomposition: The authors prove that the Expectation-Maximization (EM) objective for PRISM decomposes exactly into independent subproblems, requiring no variational approximation.
     • E-Step: Computes the posterior responsibility $P(z_t=k | \xi, \psi, \Theta)$ for each timestep. Due to the factorization of the posterior, this step scales linearly as $O(nK)$ per trajectory, a significant improvement over the $O(nK^2)$ cost of Markov-chain-based methods.
     • M-Step:
       • Reward Update: Each reward function $r_k$ is updated independently by solving a weighted Inverse Action-Value Iteration (IAVI) problem in closed form.
       • Network Update: The gating network parameters $\theta$ are updated via gradient descent to maximize the likelihood of the assigned intentions, regularized by $\ell_1$ (to encourage sparse switching) and KL-divergence (to encourage smooth transitions).

3. Training Objective:
   The loss function combines the negative log-likelihood of the intention assignments with temporal smoothness penalties to prevent erratic switching while preserving the ability to switch when necessary.

Key Contributions

1. PRISM Framework: A novel multi-intention IRL method that uses a recurrent network to model intention dynamics, enabling end-to-end learning of history-dependent switching without manual state augmentation.
2. Theoretical Guarantees: A proof that the EM objective decomposes exactly into independent per-intention reward subproblems, solvable in closed form via IAVI.
3. Computational Efficiency: An $O(nK)$ E-step achieved through posterior factorization, avoiding the quadratic complexity of forward-backward algorithms used in Markov-based approaches.
4. First Large-Scale Robotic Application: The application of multi-intention IRL to the BridgeData V2 robotic manipulation dataset, demonstrating the method's ability to generalize beyond neuroscience.
5. Interpretability: The recovery of nameable, temporally coherent intentions and reward maps from unlabeled demonstrations across three distinct domains.

Experimental Results
PRISM was evaluated on three datasets of increasing complexity:

1. Frustration Gridworld (Non-Markovian Validation):

   • Setup: A controlled environment where an agent switches strategies based on a hidden "frustration counter" (a strictly non-Markovian state).
   • Result: PRISM achieved the highest test log-likelihood and lowest expected value difference compared to HIQL, IAVI, and MaxEnt IRL. Crucially, PRISM successfully tracked the accumulating frustration and switched intentions sharply, whereas HIQL (memoryless) failed to capture the cumulative nature of the switch.

2. Mouse Labyrinth (Biological Data):

   • Setup: Real-world data of water-restricted mice navigating a 127-node labyrinth.
   • Result: PRISM outperformed SWIRL and other baselines in test log-likelihood. It recovered three distinct, interpretable intentions: water-seeking, homing, and exploration. The recovered reward maps and segmentation boundaries aligned with known biological drives and behavioral events (e.g., water-port visits).

3. BridgeData V2 (Robotic Manipulation):

   • Setup: A large-scale dataset of human-teleoperated robotic manipulation tasks with high-dimensional visual observations.
   • Result: Without any supervision or prior knowledge of intentions, PRISM discovered four temporally coherent manipulation phases: Approach/Depart, Grasp, Carry, and Idle. The method achieved higher held-out log-likelihood than baselines, demonstrating that discrete intention switching is a general property of goal-directed behavior in both biological and artificial agents.

Significance and Claims
The paper claims that PRISM provides a scalable, theoretically grounded solution for recovering latent goals in sequential behavior. By proving that the EM objective decomposes exactly, the authors show that complex history-dependent switching can be modeled without sacrificing the computational tractability of closed-form reward recovery.

The significance lies in the demonstration that:

• Discrete goal switching is prevalent in both biological and artificial agents.
• History-dependent switching (non-Markovian dynamics) can be learned end-to-end without manual feature engineering or state augmentation.
• The recovered reward maps offer a structured, interpretable lens for understanding the latent objectives behind complex demonstrations, moving beyond simple policy imitation to understanding why an agent acts at specific moments.

The authors acknowledge limitations, including the current reliance on discrete MDPs (requiring discretization for visual data) and offline operation. They suggest future work could explore model-free extensions (replacing IAVI with model-free inverse Q-learning) and online extensions for real-time intention monitoring.