Repository: pytorch/rl

Issue:
[BUG] StepCounter overrides native truncation.
## Describe the bug
When one adds a step counter on a Gymnasium env with a max_step higher than the env's native limit (TimeLimit), the env is truncated at the native TimeLimit but the truncation key is overridden to False as not truncated by the step counter.

## To Reproduce

```python
from torchrl.collectors import SyncDataCollector
from torchrl.envs import StepCounter, TransformedEnv
from torchrl.envs.libs.gym import GymEnv

import gymnasium


N = 2000
env_id = "Acrobot-v1"
if __name__ == "__main__":
    env = GymEnv(env_id)
    env = TransformedEnv(env)
    # env.append_transform(StepCounter(max_steps=N, truncated_key="truncated"))

    collector = SyncDataCollector(env, policy=None, frames_per_batch=N + 10, total_frames=2 * N)
    b = collector.next()

    print(f'done: {b["next", "done"].squeeze(-1).nonzero().squeeze(-1)}')
    print(f'truncated: {b["next", "truncated"].squeeze(-1).nonzero().squeeze(-1)}')
    print(f'terminated: {b["next", "terminated"].squeeze(-1).nonzero().squeeze(-1)}')
```

Gives

```
done: tensor([ 499,  999, 1499, 1999])
truncated: tensor([ 499,  999, 1499, 1999])
terminated: tensor([], dtype=torch.int64)
```

Adding the step counter gives:

```
done: tensor([ 499,  999, 1499, 1999])
truncated: tensor([], dtype=torch.int64)
terminated: tensor([], dtype=torch.int64)
```

## Expected behavior

The StepCounter should read if a truncation key with the same name already exists.


## System info

```bash
torch                      2.2.0a0+81ea7a4
torchrl-nightly            2024.1.19
tensordict-nightly         2024.1.19
gymnasium                  0.29.1
```


Discussion:
I think `env.env` should be `env._env`

if you do `env.env` GymEnv does not find it and falls back on your gym wrapper, gets its env (which is the order enforcing wrapper).

That is for the first part.

For step counter we should have something like `truncated = truncated | new_truncated`, but I think we simply override here. I can patch that easily I think.
Ah yes! It's `._env` not `.env`! I'm updating the issue.
Only the override is left then. I think it would be good not to override. What do you think?
One can argue that you should give a different truncation key to the StepCounter, but I wouldn't agree as both have the same semantic (why would you have 2 truncation keys?)

Resolving patch:
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
--- a/torchrl/envs/transforms/transforms.py
+++ b/torchrl/envs/transforms/transforms.py
@@ -5169,6 +5169,7 @@ def _reset(
             tensordict_reset.set(step_count_key, step_count)
             if self.max_steps is not None:
                 truncated = step_count >= self.max_steps
+                truncated = truncated | tensordict_reset.get(truncated_key, False)
                 if self.update_done:
                     # we assume no done after reset
                     tensordict_reset.set(done_key, truncated)
@@ -5187,8 +5188,10 @@ def _step(
             step_count = tensordict.get(step_count_key)
             next_step_count = step_count + 1
             next_tensordict.set(step_count_key, next_step_count)
+
             if self.max_steps is not None:
                 truncated = next_step_count >= self.max_steps
+                truncated = truncated | next_tensordict.get(truncated_key, False)
                 if self.update_done:
                     done = next_tensordict.get(done_key, None)
                     terminated = next_tensordict.get(terminated_key, None)


Test patch:
diff --git a/test/test_transforms.py b/test/test_transforms.py
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -1477,6 +1477,22 @@ def test_step_count_gym(self):
         env.rollout(1000)
         check_env_specs(env)
 
+    @pytest.mark.skipif(not _has_gym, reason="no gym detected")
+    def test_step_count_gym_doublecount(self):
+        # tests that 2 truncations can be used together
+        env = TransformedEnv(
+            GymEnv(PENDULUM_VERSIONED),
+            Compose(
+                StepCounter(max_steps=2),
+                StepCounter(max_steps=3),  # this one will be ignored
+            ),
+        )
+        r = env.rollout(10, break_when_any_done=False)
+        assert (
+            r.get(("next", "truncated")).squeeze().nonzero().squeeze(-1)
+            == torch.arange(1, 10, 2)
+        ).all()
+
     @pytest.mark.skipif(not _has_dm_control, reason="no dm_control detected")
     def test_step_count_dmc(self):
         env = TransformedEnv(DMControlEnv("cheetah", "run"), StepCounter(max_steps=30))