I have made a fork and made my changes there. I also asked for feedback on this forum a couple of days ago. (link to thread: [Contribution] [Help needed] Implementing easy action masking for distributional and dueling DQN). Perhaps we can talk about it more there?