IterableDataset 87 was introduced recently, which could be a good candidate for your use case.
I’m not deeply familiar with RL, but based on the descriptions in this topic, it seems the map-style datasets have some limitations in RL setups.