Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Extend TensorDictPrimer default_value options #2071

Merged
merged 21 commits into from
Apr 18, 2024
Merged

[Feature] Extend TensorDictPrimer default_value options #2071

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
c46e0fd
extend_primer_transform
albertbou92 Apr 9, 2024
822c04c
extend_primer_transform
albertbou92 Apr 9, 2024
c5f1a47
extend_primer_transform
albertbou92 Apr 9, 2024
667e21e
extend_primer_transform
albertbou92 Apr 9, 2024
e088240
extend_primer_transform
albertbou92 Apr 9, 2024
28525cd
extend_primer_transform
albertbou92 Apr 9, 2024
9e6c415
fix test
albertbou92 Apr 9, 2024
657bf9f
fix test
albertbou92 Apr 9, 2024
25fc14a
fix test
albertbou92 Apr 9, 2024
b97f727
__repr__
albertbou92 Apr 9, 2024
79f8772
mv validated
albertbou92 Apr 9, 2024
8895b62
fix tests
albertbou92 Apr 10, 2024
e3dcfb9
minor fix
albertbou92 Apr 10, 2024
a1cb9a1
minor fix
albertbou92 Apr 14, 2024
5a86c25
suggested changes
albertbou92 Apr 14, 2024
a64a97e
feedback changes
albertbou92 Apr 16, 2024
311f40b
feedback changes
albertbou92 Apr 16, 2024
3219150
feedback changes
albertbou92 Apr 16, 2024
891bd42
fix tests
albertbou92 Apr 16, 2024
9853121
Merge remote-tracking branch 'origin/main' into extend_primer_transform
vmoens Apr 18, 2024
1625514
amend
vmoens Apr 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
fix test
  • Loading branch information
@albertbou92
albertbou92 committed Apr 9, 2024
commit 657bf9f308d6c3da6bd0dca8e17650f30b39aefa
12 changes: 9 additions & 3 deletions test/test_transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6406,11 +6406,17 @@ def test_trans_parallel_env_check(self):
finally:
env.close()

@pytest.mark.parametrize("spec_shape", [[4], [2, 4]])
def test_trans_serial_env_check(self, spec_shape):
def test_trans_serial_env_check(self):
with pytest.raises(RuntimeError, match="The leading shape of the primer specs"):
env = TransformedEnv(
SerialEnv(2, ContinuousActionVecMockEnv),
TensorDictPrimer(mykey=UnboundedContinuousTensorSpec([4])),
)
_ = env.observation_spec

env = TransformedEnv(
SerialEnv(2, ContinuousActionVecMockEnv),
TensorDictPrimer(mykey=UnboundedContinuousTensorSpec(spec_shape)),
TensorDictPrimer(mykey=UnboundedContinuousTensorSpec([2, 4])),
)
check_env_specs(env)
assert "mykey" in env.reset().keys()
Expand Down
8 changes: 1 addition & 7 deletions torchrl/envs/transforms/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4592,14 +4592,8 @@ def transform_observation_spec(
for key, spec in self.primers.items():
if spec.shape[: len(observation_spec.shape)] != observation_spec.shape:
try:
# expanded_spec = self._try_expand_shape(spec)
expanded_spec = spec
expanded_spec = self._try_expand_shape(spec)
except AttributeError:
pass
if (
expanded_spec.shape[: len(observation_spec.shape)]
!= observation_spec.shape
):
raise RuntimeError(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When will this be reached?

Copy link
Contributor Author

@albertbou92 albertbou92 Apr 14, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if for any reason self.parent is None

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

when would transform_observation_spec be called when parent is None?

f"The leading shape of the primer specs ({self.__class__}) should match the one of the "
f"parent env. Got observation_spec.shape={observation_spec.shape} but the '{key}' entry's "
Expand Down