Repository: pytorch/rl

Issue:
[BUG] `step_mdp` exclude arguments diverge from documentation
In the docstring of `step_mdp` (https://github.com/pytorch/rl/blob/main/torchrl/envs/utils.py#L53) examples are given using the following tensordict

```pyhton
td = TensorDict(
        {
            "done": torch.zeros((), dtype=torch.bool),
            "reward": torch.zeros(()),
            "extra": torch.zeros(()),
            "next": TensorDict(
                {
                    "done": torch.zeros((), dtype=torch.bool),
                    "reward": torch.zeros(()),
                    "obs": torch.zeros(()),
                },
                [],
            ),
            "obs": torch.zeros(()),
            "action": torch.zeros(()),
        },
        [],
    )
```

the results of the following snippets in main does not align with the ones in the docs and is quite counterintuitive
```python
res = step_mdp(td, exclude_done=True)
assert "done" not in res.keys() # Error
```

```python
res = step_mdp(td, exclude_reward=True)
assert "reward" not in res.keys() # Error
```

## Cause

This is caused by the fact that when `keep_other=True` (like in this cases since it is default) the root keys will be kept as is, so the flags `exclude_reward` and `exclude_done` will prevent the keys in next to be used, but `keep_other=True` will keep the root ones.

## Solutions

I see 2 possibilities

1. make the `exclude_reward` and `exclude_done` behave as in the docs and actually exclude both the ones in root and next
2. if for some some reason this is intended a refactoring of the arg names is needed


Discussion:

Resolving patch:
diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
--- a/torchrl/envs/utils.py
+++ b/torchrl/envs/utils.py
@@ -16,7 +16,7 @@
     set_interaction_mode as set_exploration_mode,
     set_interaction_type as set_exploration_type,
 )
-from tensordict.tensordict import LazyStackedTensorDict, TensorDictBase
+from tensordict.tensordict import LazyStackedTensorDict, NestedKey, TensorDictBase
 
 __all__ = [
     "exploration_mode",
@@ -49,11 +49,14 @@ def step_mdp(
     exclude_reward: bool = True,
     exclude_done: bool = False,
     exclude_action: bool = True,
+    reward_key: NestedKey = "reward",
+    done_key: NestedKey = "done",
+    action_key: NestedKey = "action",
 ) -> TensorDictBase:
     """Creates a new tensordict that reflects a step in time of the input tensordict.
 
     Given a tensordict retrieved after a step, returns the :obj:`"next"` indexed-tensordict.
-    THe arguments allow for a precise control over what should be kept and what
+    The arguments allow for a precise control over what should be kept and what
     should be copied from the ``"next"`` entry. The default behaviour is:
     move the observation entries, reward and done states to the root, exclude
     the current action and keep all extra keys (non-action, non-done, non-reward).
@@ -76,6 +79,12 @@ def step_mdp(
             be kept in the root tensordict (since it should not be present in
             the ``"next"`` entry).
             Default is ``True``.
+        reward_key (key, optional): the key where the reward is written. Defaults
+            to "reward".
+        done_key (key, optional): the key where the done is written. Defaults
+            to "done".
+        action_key (key, optional): the key where the action is written. Defaults
+            to "action".
 
     Returns:
          A new tensordict (or next_tensordict) containing the tensors of the t+1 step.
@@ -83,6 +92,7 @@ def step_mdp(
     Examples:
     This funtion allows for this kind of loop to be used:
         >>> from tensordict import TensorDict
+        >>> import torch
         >>> td = TensorDict({
         ...     "done": torch.zeros((), dtype=torch.bool),
         ...     "reward": torch.zeros(()),
@@ -100,8 +110,7 @@ def step_mdp(
             fields={
                 done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
                 extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                reward: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
+                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
             batch_size=torch.Size([]),
             device=None,
             is_shared=False)
@@ -109,17 +118,17 @@ def step_mdp(
         TensorDict(
             fields={
                 extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                reward: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
+                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
             batch_size=torch.Size([]),
             device=None,
             is_shared=False)
-        >>> print(step_mdp(td, exclude_reward=True))  # "reward" is dropped
+        >>> print(step_mdp(td, exclude_reward=False))  # "reward" is kept
         TensorDict(
             fields={
                 done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
                 extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
+                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
+                reward: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
             batch_size=torch.Size([]),
             device=None,
             is_shared=False)
@@ -129,8 +138,7 @@ def step_mdp(
                 action: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                 done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
                 extra: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                reward: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
+                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
             batch_size=torch.Size([]),
             device=None,
             is_shared=False)
@@ -138,8 +146,7 @@ def step_mdp(
         TensorDict(
             fields={
                 done: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.bool, is_shared=False),
-                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
-                reward: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
+                obs: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
             batch_size=torch.Size([]),
             device=None,
             is_shared=False)
@@ -159,6 +166,9 @@ def step_mdp(
                     exclude_reward=exclude_reward,
                     exclude_done=exclude_done,
                     exclude_action=exclude_action,
+                    reward_key=reward_key,
+                    done_key=done_key,
+                    action_key=action_key,
                 )
                 for td, ntd in zip(tensordict.tensordicts, next_tensordicts)
             ],
@@ -169,32 +179,21 @@ def step_mdp(
             return next_tensordict
         return out
     out = tensordict.get("next").clone(False)
-    excluded = None
-    if exclude_done and exclude_reward:
-        excluded = {"done", "reward"}
-    elif exclude_reward:
-        excluded = {"reward"}
-    elif exclude_done:
-        excluded = {"done"}
-    if excluded:
+    excluded = set()
+    if exclude_done:
+        excluded.add(done_key)
+    if exclude_reward:
+        excluded.add(reward_key)
+    if len(excluded):
         out = out.exclude(*excluded, inplace=True)
-    # TODO: make it work with LazyStackedTensorDict
-    # def _valid_key(key):
-    #     if key == "next" or key in out.keys():
-    #         return False
-    #     if exclude_action and key == "action":
-    #         return False
-    #     if keep_other or key == "action":
-    #         return True
-    #     return False
     td_keys = None
     if keep_other:
         out_keys = set(out.keys())
-        td_keys = set(tensordict.keys()) - out_keys - {"next"}
+        td_keys = set(tensordict.keys()) - out_keys - excluded - {"next"}
         if exclude_action:
-            td_keys = td_keys - {"action"}
+            td_keys = td_keys - {action_key}
     elif not exclude_action:
-        td_keys = {"action"}
+        td_keys = {action_key}
 
     if td_keys:
         # update does some checks that we can spare


Test patch:
diff --git a/test/test_env.py b/test/test_env.py
--- a/test/test_env.py
+++ b/test/test_env.py
@@ -1021,6 +1021,8 @@ def test_steptensordict(
     torch.manual_seed(0)
     tensordict = TensorDict(
         {
+            "reward": torch.randn(4, 1),
+            "done": torch.zeros(4, 1, dtype=torch.bool),
             "ledzep": torch.randn(4, 2),
             "next": {
                 "ledzep": torch.randn(4, 2),