The State object itself is unnecessarily limited to (observation, reward, done, mask), etc. It should be possible to construct a State with arbitrary keys. Additionally, the state should be nestable, in it should be able to contain dictionaries or keys whose leaf nodes are eventually converted to tensors. Perhaps State should even be a subclass of a more general object.