[Feature Request] Pass in the carry tensordict upon reset in collectors

Motivation

Occasionally, we want access to what's inside the tensordict upon reset. However, with the current implementation, Env._reset would receive a tensordict with only _reset:

def _step_and_maybe_reset():
  ...
  if done_or_terminated.any():
    _reset = done_or_terminated
    td_reset = self._tensordict.select().set("_reset", _reset)
    td_reset = self.env.reset(td_reset)
    td_reset.del_("_reset")

Use case 1: Curriculum Learning.

We want to sample from some initial state distribution whose parameters are specified in the input tensordict.

Use case 2: Recording History.

It is common in robotics to use a 1-d convolution on a stacked history to perform adaptation. A graceful way to obtain the history observation is through a Transform which roughly looks like this:

class History(Transform):
    def __init__(
        self,
        in_keys: Sequence[str],
        out_keys: Sequence[str]=None,
        steps: int = 32,
    ):
        if out_keys is None:
            out_keys = [
                f"{key}_h" if isinstance(key, str) else key[:-1] + (f"{key[-1]}_h",)
                for key in in_keys
            ]
        if any(key in in_keys for key in out_keys):
            raise ValueError
        super().__init__(in_keys=in_keys, out_keys=out_keys)
        self.steps = steps
    
    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
        for in_key, out_key in zip(self.in_keys, self.out_keys):
            is_tuple = isinstance(in_key, tuple)
            if in_key in observation_spec.keys(include_nested=is_tuple):
                spec = observation_spec[in_key]
                spec = spec.unsqueeze(-1).expand(*spec.shape, self.steps)
                observation_spec[out_key] = spec
        return observation_spec

    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
        for in_key, out_key in zip(self.in_keys, self.out_keys):
            item = tensordict.get(in_key)
            item_history = tensordict.get(out_key)
            item_history[..., :-1] = item_history[..., 1:]
            item_history[..., -1] = item
        return tensordict

    def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
        for in_key, out_key in zip(self.in_keys, self.out_keys):
            item = tensordict.get(in_key)
            item_history = tensordict.get(out_key).clone()
            item_history[..., :-1] = item_history[..., 1:]
            item_history[..., -1] = item
            tensordict.set(("next", out_key), item_history)
        return tensordict

    def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
        _reset = tensordict.get("_reset", None)
        if _reset is None:
            _reset = torch.ones(tensordict.batch_size, dtype=bool, device=tensordict.device)
        for in_key, out_key in zip(self.in_keys, self.out_keys):
            if out_key not in tensordict.keys(True, True):
                item = tensordict.get(in_key)
                item_history = (
                    item.unsqueeze(-1)
                    .expand(*item.shape, self.steps)
                    .clone()
                    .zero_()
                )
                tensordict.set(out_key, item_history)
            else:
                item_history = tensordict.get(out_key)
                item_history[_reset] = 0.
        return tensordict

In reset, since we generally want the Transform itself to be stateless, it needs access to the history carried in the tensordict.

Solution

I am probably missing something. But why not just pass in the carry tensordict, i.e., let td_reset = self._tensordict.set("_reset", _reset)?

Additional context

Recoding history seems a generally useful and desired functionality. Is there a better or recommended solution?

Checklist

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