[Feature Request] Partial steps in env

[vmoens]

Motivation

#2355 would be much cleaner if we could do partial steps in batched or stateless envs.

Design question

• Should we index the batched env to make a partial step?
• Should we pass an entry like "_reset" (say, "_step") that indicates who's to be stepped?