[Feature] Add Step Counter transform

[riiswa]

Description

Add a transform that counts the steps from a reset and sets the done state to True after a certain number of steps.

This transform can be used in parallel / multi env settings, as the counter has the size of the env.batch_size. The "step_counter" key will be added to the output of tensordict when reset or step in the environment.

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• New feature (non-breaking change which adds core functionality)
• Documentation (update in the documentation)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[facebook-github-bot]

Hi @riiswa!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[riiswa]

@vmoens I'm not sure I fully understood how to edit the observation_spec, in any case here is my PR, I'm waiting for your changes suggestions :)

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[codecov]

Codecov Report

Merging #756 (922d61a) into main (f6df86c) will increase coverage by 0.02%.
The diff coverage is 96.42%.

@@            Coverage Diff             @@
##             main     #756      +/-   ##
==========================================
+ Coverage   88.72%   88.74%   +0.02%     
==========================================
  Files         123      123              
  Lines       20944    21000      +56     
==========================================
+ Hits        18583    18637      +54     
- Misses       2361     2363       +2     

Flag	Coverage Δ
habitat-gpu	24.86% <29.16%> (+0.01%)	⬆️
linux-brax	29.49% <29.16%> (-0.01%)	⬇️
linux-cpu	85.28% <96.42%> (+0.03%)	⬆️
linux-gpu	86.18% <96.42%> (+0.01%)	⬆️
linux-jumanji	30.27% <29.16%> (-0.01%)	⬇️
linux-outdeps-gpu	72.14% <96.42%> (+0.26%)	⬆️
linux-stable-cpu	85.14% <96.42%> (+0.04%)	⬆️
linux-stable-gpu	85.84% <96.42%> (+0.02%)	⬆️
linux_examples-gpu	42.88% <29.16%> (-0.03%)	⬇️
macos-cpu	85.04% <96.42%> (+0.10%)	⬆️
olddeps-gpu	76.00% <96.42%> (+0.12%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
torchrl/envs/__init__.py	100.00% <ø> (ø)
torchrl/envs/transforms/__init__.py	100.00% <ø> (ø)
torchrl/envs/transforms/transforms.py	87.38% <91.66%> (+0.08%)	⬆️
test/test_transforms.py	96.64% <100.00%> (+0.08%)	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[vmoens]

Good work!
Have a look at my comments

[vmoens]

Maybe i'm missing something, but I think you should do something like

    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
        observation_spec["step_count"] = UnboundedDiscreteTensorSpec(dtype=torch.int64)

[riiswa]

Thanks for all your suggestions! I made the changes, but I would like to know how can I test the transform_observation_spec method

[vmoens]

There's a test function you can run

from torchrl.envs.utils import check_env_specs
env = TransformedEnv(GymEnv("Pendulum-v1"), StepCounter(50))
check_env_specs(env)

It will run a small rollout + get fake data from your specs, and check that they match.

[riiswa]

I get this error, I don't really know what I missed, it comes from the shape of the tensor is []:

test/test_transforms.py:1600 (TestTransforms.test_step_counter_observation_spec)
Traceback (most recent call last):
  File "/Users/waris/Projects/rl/test/test_transforms.py", line 1604, in test_step_counter_observation_spec
    check_env_specs(env)
  File "/Users/waris/Projects/rl/torchrl/envs/utils.py", line 168, in check_env_specs
    fake_tensordict = env.fake_tensordict().flatten_keys(".")
  File "/Users/waris/Projects/rl/torchrl/envs/common.py", line 636, in fake_tensordict
    fake_obs = observation_spec.rand(self.batch_size)
  File "/Users/waris/Projects/rl/torchrl/data/tensor_specs.py", line 1286, in rand
    _dict = {
  File "/Users/waris/Projects/rl/torchrl/data/tensor_specs.py", line 1287, in <dictcomp>
    key: self[key].rand(shape)
  File "/Users/waris/Projects/rl/torchrl/data/tensor_specs.py", line 655, in rand
    r = torch.rand(*shape, *interval.shape, device=interval.device)
TypeError: rand() missing 1 required positional arguments: "size"

[vmoens]

On line 655, try putting the shape in between brackets
Like rand([*shape1, *shape2],...)
Also make sure you've merged main into your branch!

[vmoens]

LGTM
This should be adapted to the Nd*Spec removal though :)