$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 89. Improved: $\large\color{#35bf28}5$. Worsened: $\large\color{#d91a1a}3$.

Putting this on hold until #1448 has landed
THe _step will be drastically simplified and as a result I expect that the behaviour of TensorDictPrimer will be clearer too.

@albertbou92 tests are now passing but i'm still not 100% confident about TensorDictPrimer.

First let me explain what TensorDictPrimer is about:

• Collectors allocate buffers to communicate from process to process based on the environment specs. We basically do env.fake_tensordict() + some shape expansion and use that as a data container for communication.
• If there is some data that needs to be carried but is not produced by the env but by the policy, we need to pass it too. Therefore, we designed TensorDictPrimer, which adds the missing tensors to the env specs such that env.fake_tensordict is complete.
• In our case, we want to add the input and output hidden states to the specs.
• The question is whether we want these to be added to the input or output spec. By definition, output is what the env is producing so that would not make much sense to put it there (in the context of LSTMs). Therefore, we should probably put it in input. However, when calling reset in batched environments, we filter out the data that is not part of output_spec. This means that when calling reset, we won't see the preallocated tensors from TensorDictPrimer. One possible fix would be to also keep the tensors from input_spec if they can be found.
• The contract of TensorDictPrimer is nevertheless that the values will be available after reset, and as such this should be part of output_spec too.

My main concern is basically that this is all a bit hacky: we're telling the env that it must expect tensors as input when, in fact, they have nothing to do with the env. Perhaps we could add a spec to the policy instead and tell collectors to fetch the specs from the policy as well as the env, and combine the two. That way we'd keep things separated.

One long-term solution I could propose is to have SafeModules that also have input and output specs.
The advantages are:

• Clear separation between data origin (policy vs env)
• transforms can be moved from env to policy without problems
• We will be able to call parent or container on a transform stored within the policy.

There are drawbacks:

• Currently SafeModules are only accessory (they just do some more checks than regular modules). This would force users to use SafeModules everywhere, which may not be the best experience.
• The parent and container interface won't be very clean in the transforms.
• It could be error prone (transforming the env output specs in policy input specs may not be straightforward). Overall, I would expect that the user experience will be negatively impacted, which is a high price to pay. For instance, it could be something like

env = make_env()
policy = SafeActor(network, input_spec=CompositeSpec(obs=...), output_spec=CompositeSpec(action=...))

@smorad @matteobettini wdyt?

Safe modules are a nice optional feature, but forcing practitioners to use them is not very user-friendly.

Regarding the LSTM hidden states, the TensorDictPrimer transform places them in the TensorDict and you already obtain them after calling reset. During the reset, the environment creates these hidden states, and in every subsequent step, it passes them forward without changing them but still returns them. So, in my view, it's not unreasonable to consider them as outputs.

In fact, same way you can think of output as what the environment produces, you can think of input as what the environment requires to take a step. However, not all the tensors we define in the input_spec fit this definition. For example, when using the RewardSum transform, the episode_reward specs are part of the input_specs, even though they are not something the environment uses to perform steps.

Wouldn't forcing every modules to be Safe in the library be quite bc-breaking?
IMO the impact of that solution might be substantial.

It seems like keeping the primers in output spec is the best solution as of now. at the end of the day they are still data that needs to be given to step_mdp and as policy input