[Feature Request] Metadata for specs

[vmoens]

          > Hi there. Besides the naming, what do you think of adding some metadata for users to populate?

This could be useful for, for example, marking which keys are visual observations or temporally stacked low-dimensional observations. People commonly differentiate the type of obs through their dimensions but in this case, they have the same number of dimensions but are expected to be handled differently.

It is in general useful for subsequent algorithm implementation to parse the observations and build the network.

Like this:

visual_obs_spec = Continuous(shape=[N, W, H])
visual_obs_spec.meta["is_pixel"] = True
stacked_obs_spec = Continuous(shape=[N, T, D])
stacked_obs_spec.meta["is_pixel"] = False

@btx0424

I like the idea but do we want more specific flags? We could think of an API where there is a list of metadata you can put:

class SpecMetadata(Enum):
    PIXELS = 0
    STATE = 1
    ACTION = 2
    ...

visual_obs_spec = Unbounded((C, H, W), dtype=torch.uint8)
state_obs_spec = Unbounded((F,), dtype=torch.float32)

visual_obs_spec.metadata.add(SpecMetadata.PIXELS)

state_obs_spec.metadata.add(SpecMetadata.STATE)

visual_obs_spec.check_metadata(SpecMetadata.STATE) # False
visual_obs_spec.check_metadata(SpecMetadata.PIXELS) # True

Originally posted by @vmoens in https://github.com/pytorch/rl/issues/1818#issuecomment-2271322059

[btx0424]

Thanks for your timely response!

As research nowadays is shifting to more complex and realistic environments, we might need more and more complex metadata. I personally prefer more flexibility, e.g., just restricting the custom metadata to be of primitive types for serialization purposes.

[vmoens]

My idea was for this to be extensible at will, but using strings makes it clunky (it's easy to put a cap letter or a typo)