Learning to Undo: Rollback-Augmented Reinforcement Learning with Reversibility Signals
Andrejs Sorstkins, Omer Tariq, Muhammad Bilal
TL;DR
This paper addresses safety and learning stability in value-based reinforcement learning by introducing a reversible framework that jointly learns an online per-state–action reversibility signal, $\,\Phi(s,a)\$, and a rollback operator. The reversibility signal quantifies the probability of returning to a prior state within a horizon $K$, and it modulates TD targets through a penalty $r' = r - \lambda(1-\Phi[s_t,a_t])$ to reduce overestimation-driven instability. A selective rollback triggers when a target underperforms a threshold, forcing the agent to revert to the previous state (and action in SARSA), effectively interrupting high-risk trajectories and preventing catastrophic updates. Empirical results in CliffWalking-v0 and Taxi-v3 show dramatic safety improvements (e.g., reductions of catastrophic falls and illegal actions by >99%), substantial gains in cumulative rewards (≈55–66%), and pronounced reduction in reward variance, with ablations confirming rollback as the key contributor. The framework demonstrates environment-adaptive parameter profiles and highlights the practical potential of reversibility and rollback as a safe, robust building block for sequential decision making, with clear avenues for extending to deep RL and real-world robotics.
Abstract
This paper proposes a reversible learning framework to improve the robustness and efficiency of value based Reinforcement Learning agents, addressing vulnerability to value overestimation and instability in partially irreversible environments. The framework has two complementary core mechanisms: an empirically derived transition reversibility measure called Phi of s and a, and a selective state rollback operation. We introduce an online per state action estimator called Phi that quantifies the likelihood of returning to a prior state within a fixed horizon K. This measure is used to adjust the penalty term during temporal difference updates dynamically, integrating reversibility awareness directly into the value function. The system also includes a selective rollback operator. When an action yields an expected return markedly lower than its instantaneous estimated value and violates a predefined threshold, the agent is penalized and returns to the preceding state rather than progressing. This interrupts sub optimal high risk trajectories and avoids catastrophic steps. By combining reversibility aware evaluation with targeted rollback, the method improves safety, performance, and stability. In the CliffWalking v0 domain, the framework reduced catastrophic falls by over 99.8 percent and yielded a 55 percent increase in mean episode return. In the Taxi v3 domain, it suppressed illegal actions by greater than or equal to 99.9 percent and achieved a 65.7 percent improvement in cumulative reward, while also sharply reducing reward variance in both environments. Ablation studies confirm that the rollback mechanism is the critical component underlying these safety and performance gains, marking a robust step toward safe and reliable sequential decision making.