Masked Skill Token Training for Hierarchical Off-Dynamics Transfer

Zeyu Feng¹, Haiyan Yin¹, Yew-Soon Ong^1,3, Harold Soh^2,4

¹CFAR and IHPC, Agency for Science, Technology and Research (A*STAR), Singapore ²Department of Computer Science, School of Computing, National University of Singapore (NUS), Singapore ³Nanyang Technological University (NTU), Singapore ⁴Smart Systems Institute, NUS

This repository contains the implementation of the masked training method for skill transfer proposed in the ICLR 2026 paper Masked Skill Token Training for Hierarchical Off-Dynamics Transfer.

📢 Introduction

Generalizing policies across environments with altered dynamics remains a key challenge in reinforcement learning, particularly in offline settings where direct interaction or fine-tuning is impractical. We introduce Masked Skill Token Training (MSTT), a fully offline hierarchical RL framework that enables policy transfer using observation-only demonstrations. MSTT constructs a discrete skill space via unsupervised trajectory tokenization and trains a skill-conditioned value function using masked Bellman updates, which simulate dynamics shifts by selectively disabling skills. A diffusion-based trajectory generator, paired with feasibility-based filtering, enables the agent to execute valid, temporally extended actions without requiring action labels or access to the target environment. Our results in both discrete and continuous domains demonstrate the potential of mask-guided planning for robust generalization under dynamics shifts. To our knowledge, MSTT is the first work to explore masking as a mechanism for simulating and generalizing across off-dynamics environments. It marks a promising step toward scalable, structure-aware transfer and opens avenues to explore multi-goal conditioning, and extensions to more complex, real-world scenarios.

🔨 Usage

We will release the code for training and testing in this repository in the near future.

🏷️ License

This repository is released under the MIT license. See LICENSE for additional details.

If you find this repo or the ideas presented in our paper useful for your research, please consider citing our paper.

@inproceedings{feng2026masked,
  title     = {Masked Skill Token Training for Hierarchical Off-Dynamics Transfer},
  author    = {Zeyu Feng and Haiyan Yin and Yew-Soon Ong and Harold Soh},
  booktitle = {The Fourteenth International Conference on Learning Representations},
  year      = {2026},
  url       = {https://openreview.net/forum?id=K4ngUOra9m}
}

🙏 Acknowledgement

• Our implementations are based on Diffuser and DOPPLER.