Request for Example of AEC API Usage with Agent Masking in Petting Zoo

I've been exploring the BenchMARL library and am impressed with its capabilities and design—great work!

I am currently interested in implementing a multi-agent reinforcement learning scenario using the AEC (Agent-Environment Cycle) API in petting zoo, particularly for environments that require sequential turn-based actions like in a Chess game. In this context, I need to apply masking at the agent level rather than action masking.

Could you provide an example or guidance on how to adapt the AEC API for such a use case? Any examples of AEC API usage with agent masking in a Chess-like environment would be incredibly helpful.

Thank you for your assistance and for the excellent work on BenchMARL.