Learning to Undo: Transfer Reinforcement Learning under State Space Transformations

ICLR 2026 Conference Submission21846 Authors

Keywords: reinforcement learning, transfer learning

TL;DR: We show how learning an undo map between MDPs enables efficient policy transfer under state space transformations.

Abstract: Transfer learning in reinforcement learning (RL) has shown strong empirical success. In this work, we take a more principled perspective by studying when and how transferring knowledge between MDPs can be provably beneficial. Specifically, we consider the case where there exists an undo map between two MDPs (a source and a target), such that applying this map to the target’s state space recovers the source exactly. We propose an algorithm that learns this map via regression on state feature statistics gathered from both MDPs, and then uses it to obtain the target policy in a zero-shot manner from the source policy. We theoretically justify the algorithm by analyzing the setting when the undo map is linear and the source is linearly-$Q^\star$ realizable, where our approach has strictly better sample complexity than learning from scratch. Empirically, we demonstrate that these benefits extend beyond this regime: on challenging continuous control tasks, our method achieves significantly better sample efficiency. Overall, our results highlight how shared structure between tasks can be leveraged to make learning more efficient.

Primary Area: reinforcement learning

Submission Number: 21846