[Feature Request] Refactoring of TensorSpec and documentation

[DavideTr8]

Motivation

When creating a custom environment, deciding the correct TensorSpec (TS) to use is not straightforward. Some TS are never used, some look very similar to each other and sometimes docstrings are too concise or they miss some parameters.

Solution

After a brief talk with @vmoens and @fedebotu, we propose:

1. Improving the documentation: better detail docstrings and maybe add a tutorial for explaining what TS to use for different use cases.
2. Refactoring the tensor_specs module: review all the classes, their naming, their methods and check if there is a more intuitive way to achieve the same result.

Alternatives

The refactoring might happen as a v2 of the current tensor_specs module, to avoid breaking compatibility with previous versions of the code.

Additional context

This feature is proposed as a mean to improve user experience and as a possible follow up to #1795

Checklist

• [x] I have checked that there is no similar issue in the repo (required)

[vmoens]

Sorry for totally dropping the ball on this.

Here's the long term plan:

1. Better doc: we need to make it super clear when and where each spec should be used. Right now the docstrings are almost inexistent, yet we tell people who want to design their algo that they should use them.
2. Fix issues: there are many outstanding issues right now, such as lack of clarity between what is a discrete spec (resp to a continuous spec with integer domain for instance) and similar.
3. Naming issues: The names are horrific. Long, clunky and impossible to remember. I suggest to make a shorter version of each spec (UnboundedContinuousTensorSpec -> Unbounded, etc). We could simply symlink between the old names and the new ones with a warning coming with the old names saying that these will be deprecated.

I expect (1) and (3) to live in a single PR each. (3) should be super easy to come with, (1) requires a bit of digging but if someone wants to give it a shot it'd be awesome. (2) will require more discussion and precise identification of the outstanding issues.

Eventually we could consider a specs.v2 if we see that the current API is really broken, but I'd rather not do it if we can avoid that.

Inviting @matteobettini @skandermoalla @albertbou92 @BY571 @btx0424 to the party if you're interested to comment or help (you never know 😝)

[DavideTr8]

In few days I can start working on point 1 if it's not too late. Shouldn't take too long!

[vmoens]

Renaming specs

The specs name are clunky and even I don't remember them all. We should rename them:

OneHotDiscreteTensorSpec -> OneHot(Spec)
CompositeSpec -> Composite(Spec)
UnboundedContinuousTensorSpec -> Continuous(Spec)
UnboundedDiscreteTensorSpec -> Discrete(Spec)
MultiOneHotDiscreteTensorSpec -> MultiOneHot(Spec)
NonTensorSpec -> NonTensor(Spec) or Object(Spec)
MultiDiscreteTensorSpec -> MultiDiscrete(Spec)
LazyStackedTensorSpec -> Stacked(Spec)
LazyStackedCompositeSpec -> StackedComposite(Spec)
DiscreteTensorSpec -> Categorical(Spec)
BoundedTensorSpec -> Bounded(Spec) or Finite(Spec)
BinaryDiscreteTensorSpec -> Binary(Spec)

I put the (Spec) because I'm not even sure we need it. Opinions?

[matteobettini]

Very needed I think and I love it!

What I think we need to nail is to keep namings that are simple yet mutually exclusive.

Does Bounded imply continuous? cause otherwise also the categoricals and multi hot ones are bounded. For me it would be clearer if it was called ContinuousBounded.

What is the use of Discrete ? Is it for spec of ints without a domain range? If so it should be clearer how it differs from the categoricals and one hots, otherwise it seems a super-class. Maybe have DiscreteUnbounded and ContinuousUnbounded. Eventually we could also merge ContinuousUnbounded and ContinuousBounded into Continuous with a setup in init. (The same we would not do for the discretes because as we see here we have many types of bounding for discrete)

The rest I really like

so my proposal is

CompositeSpec -> Composite

UnboundedContinuousTensorSpec -> ContinuousUnbounded
BoundedTensorSpec -> ContinuousBounded

UnboundedDiscreteTensorSpec -> DiscreteUnbounded
OneHotDiscreteTensorSpec -> OneHot
MultiOneHotDiscreteTensorSpec -> MultiOneHot
DiscreteTensorSpec -> Categorical
MultiDiscreteTensorSpec -> MultiDiscrete 
BinaryDiscreteTensorSpec -> Binary

LazyStackedTensorSpec -> Stacked
LazyStackedCompositeSpec -> StackedComposite
NonTensorSpec -> NonTensor

[vmoens]

UnboundedDiscreteTensorSpec -> DiscreteUnbounded

I would ditch that one altogether, it's just an Unbounded with a discrete dtype (int) no?

We can keep the old one and just specialize it based on Unbounded?

We'll need to keep the old names for a couple of releases (at least) anyway, so the thing we want to do is raise an appropriate warning to tell people how to use the new ones.

CompositeSpec -> Composite

UnboundedDiscreteTensorSpec / UnboundedContinuousTensorSpec -> Unbounded
# We can make ContinuousUnbounded and DiscreteUnbounded, both with predefined kwargs

BoundedTensorSpec -> Bounded
BinaryDiscreteTensorSpec -> Binary (specialized version of Bounded?)

OneHotDiscreteTensorSpec -> OneHot
MultiOneHotDiscreteTensorSpec -> MultiOneHot
DiscreteTensorSpec -> Categorical
MultiDiscreteTensorSpec -> MultiDiscrete 

LazyStackedTensorSpec -> Stacked
LazyStackedCompositeSpec -> StackedComposite

NonTensorSpec -> NonTensor

[vmoens]

@matteobettini

If you check this, you will see that instance checks with all classes are interchangeable, to make sure that people can smoothly switch from one class to the other.

For Continuous and Discrete, we refer to the domain/dtype to determine the class.

#2360

https://github.com/pytorch/rl/blob/ba582a7eea184e5423c74515712b8af85f6f35d9/test/test_specs.py#L3623-L3793

[btx0424]

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

[matteobettini]

UnboundedDiscreteTensorSpec -> DiscreteUnbounded

I would ditch that one altogether, it's just an Unbounded with a discrete dtype (int) no?

ok good i agree

We can keep the old one and just specialize it based on Unbounded?

We'll need to keep the old names for a couple of releases (at least) anyway, so the thing we want to do is raise an appropriate warning to tell people how to use the new ones.

CompositeSpec -> Composite

UnboundedDiscreteTensorSpec / UnboundedContinuousTensorSpec -> Unbounded
# We can make ContinuousUnbounded and DiscreteUnbounded, both with predefined kwargs

BoundedTensorSpec -> Bounded
BinaryDiscreteTensorSpec -> Binary (specialized version of Bounded?)

OneHotDiscreteTensorSpec -> OneHot
MultiOneHotDiscreteTensorSpec -> MultiOneHot
DiscreteTensorSpec -> Categorical
MultiDiscreteTensorSpec -> MultiDiscrete 

LazyStackedTensorSpec -> Stacked
LazyStackedCompositeSpec -> StackedComposite

NonTensorSpec -> NonTensor

My only concern now is for Bounded. The way I see it we either have a bounded continuous or, in the discrete setting, bounding means one of the binary, categorical, onehot, etc.

Therefore, just having Bounded is not super clear to me, as if i want a categorical technically it can be also a Bounded discrete.

I would let users be able to instantiate only mutually exclusive leaf specs. AKA to get a categorical they need to create a Categorical not also by having a Bounded with discrete domain

my suggestion then is

CompositeSpec -> Composite

UnboundedDiscreteTensorSpec / UnboundedContinuousTensorSpec -> Unbounded
# We can make ContinuousUnbounded and DiscreteUnbounded, both with predefined kwargs

# Bounding in continuous domain
BoundedTensorSpec -> ContinuousBounded (can also call it Interval or smth that is one word and imples both bounds and continiuity)

# Bounding in discrete domain
BinaryDiscreteTensorSpec -> Binary 
OneHotDiscreteTensorSpec -> OneHot
MultiOneHotDiscreteTensorSpec -> MultiOneHot
DiscreteTensorSpec -> Categorical
MultiDiscreteTensorSpec -> MultiDiscrete 

LazyStackedTensorSpec -> Stacked
LazyStackedCompositeSpec -> StackedComposite

NonTensorSpec -> NonTensor

[vmoens]

Therefore, just having Bounded is not super clear to me, as if i want a categorical technically it can be also a Bounded discrete.

What's the added value of having BoundedDiscrete compared to Bounded(..., dtype=torch.int).domain == "discrete"?

EDIT: I don't get the proposal, why are we having only ContinuousBounded and not DiscreteBounded?

[vmoens]

@matteobettini This is what the new doc will look like

#2368

[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

[matteobettini]

@matteobettini This is what the new doc will look like

#2368

Very cool, my main question is still what is the difference between bounded discrete, categorical and discrete?

[vmoens]

Categorical spec

Stuff you will use to index => Categorical: smth[tensor] => OneHot: smth[tensor] using masking => Binary: subclasses Categorical - but could also be seens as a OneHot where more than one element can be taken

Numerical

Unbounded: box is from -inf to inf (or better from iinfo/finfo.min to max) Bounded: Same but box is restricted to presumably non-inf values

You'd use unbounded for pixels in uint8 for instance: they're not really bounded (only by the dtype range) and they're not used to index even though they are ints. In fact, smth[pixels] would make no sense, so they can't be said to be categorical. On the other hand, an action that encodes values from 0 to 10 is not really bounded int because what you want is to enumerate things (with pixels you don't enumerate, you discretize).

[matteobettini]

I see, it is more a semantic meaning, because in practice a categorical 10 can still look like a bounded discrete 0 to 10.

What is this new Discrete type that I see in the cateogrical side of the doc?

What I suggest is run them by someone that is not familiar with torchrl (you can pick a guy from the street) and ask to guess what each class intuitively does. I think it is core that this is as much intuitive and simple as possible to users. so we need as much input as we can get.

[vmoens]

All the previous classes still exist.

The overlap between categorical, one-hot and others already existed (a one-hot tensor could pass the test to belong to a bounded / unbounded spec). I don't think anything is changing on that side. At the end of the day, these classes are not only there to specify the dtype and possible values of the tensors (otherwise we'd just need one tensor spec) but also their meaning.

In fact, we have a function to generate specs automatically (make_composite_from_td) that will be the most tolerant version of the specs you could imagine.

There is no Discrete, thanks for spotting, it should be Binary.