so, just to clarify, if I have

• an input tensor of shape [*B,T,*N,F]
• in a tensordict with shape [*B,T]

my *N dimensions are treated as batch?
Is this usecase supported?

Yes, see the test_single_step_vs_multi test

in that test the *B is part of the tensordict.

Here I am asking about the *N
like the agent dimension in MARL for example, after the T, but before the F

here we are assuming is_init has the tensordict shape [*B,T]?

not sure what makes you think that? If so is it a problem? Can you elaborate?

• What makes me think that is that we consider its last dimension as T in the lines after
• What might become a problem is that 1) we are hardcoding the name "is_init". 2) it might have a different shape as it follows a done no?

I think this issue and the one above boil down to what we want to accept as input to this module.
We need to have a way to tell what the time dimension is.

• One option is to check the dim names but it's hard to get people to adopt that. I tried and got multiple issues of people saying that their td was lacking the "time" dim name and they did not know what to do.
• Another option is to pass the time dim in the constructor but it's brittle. If you say time_dim=3 you won't be able to change the batch size, if you say time_dim=-3 you need to say with respect to what (obs, init, tensordict?)

A regular usage is assumed to be that the time is the last dim of the tensordict, but it's true that we could cover others. Otherwise we assume that the TD provided is always the root and always has time dim at the end. Handling batch-dims after the time dim can also be considered.

Another issue is that we assume that the hidden state can more or less be indexed by is_init but i'm not sure what this looks like in cases where you have multiple done states and multiple observations. The batch size of the obs could differ from the done state (eg in VMAS IIUC you have one done state per group of agents but one obs per agent).
In MARL settings you will also need to consider if the LSTM reads the aggregated observations of not, if you have one LSMT per agent or not etc. Overall it starts to look like a different class to me.

TL;DR
Happy to have a follow-up work on GRU and LSTM to make them compatible with more convoluted data structures, but do you think we can integrate GRU as it is now since it applies to most single agents use cases already?
If you have a precise list of queries of what this class should also cover i'd be glad to look at it.
What I would need is: what are the data structures that need to be considered, what behaviour we expect, and what we want to explicitly forbid because there is a solution. To give you an example, IMO the example of the done that has a shape that differs from the obs can be easily solved by forbidding that and telling users to append an rename + unsqueeze + expand transform before (which all work as regular nn.Modules).

Yes let's go ahead with this as is and we can keep discussing on this. Also because the LSTM has the same considerations that can be applied so we can tackle this in a future PR.

The example

TnesorDict{
 "is_init": Tensor(shape=[*B,T])
  "group": TensorDict{
    "obs":  Tensor(shape=[*B,T,*N,F])
    "recurrent": Tensor(shape=[*B,T,*N,H])
    batch_size=[*B, T, *N]
    }
  batch_size=[*B,T]
}

The assumptions

• i would assume that is_init, the input, and the recurrent state have all the same batch_size and the last dim respectively of 1, F, H. In the example above, is_init might have to be expanded into the group before passing it to the module
• i would assume (like in value functions) that the root td is always the one passed to the module and its last dim the time dimension. that way i know where the time dimension is (without the need to mark it) and know that all the other dimensions in my input apart from that and the last are batches

What we expect

We then woud expect that the input and recurrent state they are reshaped with batch=[*B x *N] and time=[T], ready to be passed to the inner GRU/LSTM. ands then they are reshaped back.

i would assume that is_init, the input, and the recurrent state have all the same batch_size and the last dim respectively of 1, F, H. In the example above, is_init might have to be expanded into the group before passing it to the module

my view in this case is (1) raise an exception because we can't really tell what to do anticipatively with a is_init that has a mismatching shape and (2) provide info on how to do the shape modif

If only first batch dims were working with rnns, sigh!