Exact Unlearning in Reinforcement Learning

Technical Summary: Exact Unlearning in Reinforcement Learning

Problem Formulation
The paper addresses the problem of exact unlearning in the context of Reinforcement Learning (RL). While machine unlearning has been studied in supervised and unsupervised settings, it remains largely unaddressed in interactive, sequential systems like RL. The core challenge is to design an efficient framework that allows for the removal of a specific user’s data (an episode) upon request, such that the learner’s output after unlearning is statistically indistinguishable from the output that would have been produced had that user’s data never been included in the training process.

The authors define Exact Unlearning (Definition 2.1) for sequential learning problems: An unlearning algorithm $\mathcal{U}$ achieves exact unlearning with respect to a learning algorithm $\mathcal{A}$ if, for any sequence of users $z_{1:T}$ and any target user $z_t$ to be deleted (replaced by a dummy user $z'$), the distribution of the unlearned model is identical to the distribution of the model trained on the modified sequence $z_{1:t-1}, z', z_{t+1:T}$. This is a stricter requirement than approximate unlearning or differential privacy, as it demands identical distribution rather than statistical similarity.

Methodology
The authors extend the framework of Total Variation (TV) stability from batch learning (Ullah & Arora, 2023) to sequential RL. The methodology relies on two main components:

1. TV-Stable Learning via Binary Tree Mechanism:

   • The authors abstract RL into a class of sequential learning problems with prefix sum structures.
   • They employ a binary tree mechanism (adapted from differential privacy literature) to maintain noisy prefix sums of sufficient statistics. Instead of storing raw counts, the algorithm stores clean values and noisy values in a binary tree structure. Each node in the tree holds an independent Gaussian noise sample.
   • This structure ensures that any prefix sum is perturbed by only $O(\log T)$ independent noise samples, rather than $T$ samples, which is crucial for controlling the TV distance between the original and modified datasets.
   • An algorithm is defined as $\rho$-TV-stable if the Total Variation distance between its outputs on two datasets differing by one element is bounded by $\rho$.

2. Efficient Unlearning via Maximal Coupling:

   • When a deletion request occurs, the algorithm performs unlearning by maximally coupling the binary tree of the original dataset with the tree of the modified dataset.
   • The unlearning procedure (Algorithm 3) traverses the path from the deleted leaf to the root. For each node, it attempts to reuse the existing noisy value via rejection sampling. If the rejection sampling fails (which happens with probability proportional to the TV distance), it applies a reflection map to generate a new valid noisy value and retrains only the affected subtree.
   • This approach ensures that full retraining is rare, occurring with probability $\rho \sqrt{\log T}$, making the computational cost of unlearning significantly lower than retraining from scratch.

3. RL Algorithm Design (UCB-VI Extension):

   • The authors instantiate this framework for tabular Markov Decision Processes (MDPs) by extending the UCB-VI algorithm (Azar et al., 2017).
   • They replace standard visitation statistics with noise-perturbed prefix sums.
   • To maintain optimism (a key component of UCB-VI) despite the noise, they introduce a modified bonus function that accounts for both estimation error (Hoeffding-style) and approximation error introduced by the noisy counts.

Key Contributions

1. Formulation: The paper formulates the problem of exact unlearning in RL, abstracting it into sequential learning problems with prefix sum structures.
2. Algorithm: It constructs the first regret-optimal RL algorithm for tabular MDPs that is $\rho$-TV-stable and supports efficient exact unlearning.
3. Theoretical Guarantees:
   • Regret Bound: The proposed algorithm achieves a regret bound of $\tilde{O}(H^2\sqrt{SAT} + H^3S^2A + H^{2.5}S^2A/\rho)$.
   • Unlearning Cost: The expected computational cost of unlearning is only a $\rho \sqrt{\ln T}$ fraction of the cost of retraining from scratch.
   • Lower Bound: The authors establish a minimax lower bound of $\Omega(H\sqrt{SAT} + SAH/\rho)$ for any $\rho$-TV-stable RL algorithm, showing that their upper bound is nearly tight.

Results

• Regret Analysis: The regret bound consists of three terms: a burn-in cost, the standard UCB-VI regret term (using Hoeffding bonuses), and an additional term $H^{2.5}S^2A/\rho$ accounting for the cost of stabilizing the algorithm with Gaussian noise.
• Comparison with Differential Privacy: The authors note that their framework incurs a lower additional regret cost compared to Joint Differential Privacy (JDP) RL algorithms. Specifically, their additional cost is lower by a factor of $\sqrt{H}$ because Gaussian noise provides a tighter sensitivity bound ($B=\sqrt{3H}$) compared to the Laplacian noise used in JDP ($B=3H$).
• Tightness: The derived lower bound matches the form of lower bounds for JDP algorithms but is strictly stronger because TV-stability is a weaker condition than pure DP. The gap between the upper and lower bounds is attributed to the use of Hoeffding-type bonuses instead of Bernstein-type bonuses and potential analysis gaps in the second term.

Significance
The paper claims to bridge a significant gap in the literature by providing the first theoretical and algorithmic framework for provable exact unlearning in RL. It demonstrates that it is possible to retrofit strong deletion guarantees onto regret-optimal learners while keeping retraining rare and localized. This is particularly relevant for interactive systems (e.g., recommendation systems, personalized assistants) where user data is logged episodically and deletion requests arise for privacy, safety, and compliance reasons. The work establishes TV-stability as a viable mechanism for achieving exact unlearning in sequential settings, offering a foundation for future research into unlearning in more complex RL environments (e.g., with function approximation).