Skip to content

[BUG] SelectTransform and ExcludeTransform don't Select and Exclude #1572

@pranavsetpal

Description

@pranavsetpal

The keys the inputted for the SelectTransform and ExcludeTransform do not affect the environment at all.

To Reproduce

  1. Create any environment
  2. Transform the environment using either ExcludeTransform or SelectTransform
  3. Print the observation spec before and after

Actual Behaviour
ExcludeTransform -> doesn't exclude any items
SelectTransform -> returns an empty dictionary (doesn't include any items)

from torchrl.envs.libs.vmas import VmasEnv
from torchrl.envs import TransformedEnv, ExcludeTransform

env = VmasEnv("balance", num_envs=1, n_agents=2)

print("Before Transform:", env.observation_spec.keys(True, True))

env = TransformedEnv(
    env,
    ExcludeTransform( ("agents", "info", "ground_rew") ) # Just example
)

print(" After Transform:", env.observation_spec.keys(True, True))
Before Transform: _CompositeSpecKeysView(keys=[('agents', 'observation'), ('agents', 'info', 'pos_rew'), ('agents', 'info', 'ground_rew')])
 After Transform: _CompositeSpecKeysView(keys=[('agents', 'observation'), ('agents', 'info', 'pos_rew'), ('agents', 'info', 'ground_rew')])

Expected behavior

I expected the transformed environment to exclude the keys

System info

TorchRL: 2023.9.18 (Nightly, via pip)
Python: 3.11.5

Reason and Possible fixes

I believe the issue lies in the transform_observation_spec of both classes. There the transformed spec is created

return CompositeSpec(
    {
        key: value
        for key, value in observation_spec.items()
        if unravel_key(key) not in self.excluded_keys
    },

Particularly,

        for key, value in observation_spec.items()

this doesn't include all nested keys and hence no items are being excluded from the spec.
A possible solution which works on my machine would be to replace the above line with

        for key, value in observation_spec.items(True, True)

Checklist

  • I have checked that there is no similar issue in the repo (required)
  • I have read the documentation (required)
  • I have provided a minimal working example to reproduce the bug (required)

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions