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

[BugFix] TensorDictPrimer updates spec instead of overwriting #2332

Merged
merged 8 commits into from
Jul 31, 2024
Merged

[BugFix] TensorDictPrimer updates spec instead of overwriting #2332

Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
9a4718f
amend
matteobettini Jul 29, 2024
cfd30f8
amend
matteobettini Jul 29, 2024
4e3103d
amend
matteobettini Jul 29, 2024
abdaa61
amend
matteobettini Jul 30, 2024
54a0842
Merge remote-tracking branch 'origin/main' into primer
vmoens Jul 30, 2024
25e85e1
Merge remote-tracking branch 'origin/main' into primer
vmoens Jul 30, 2024
700f0be
test
matteobettini Jul 31, 2024
50940ee
Merge branch 'refs/heads/main' into primer
matteobettini Jul 31, 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
amend
  • Loading branch information
@matteobettini
matteobettini committed Jul 29, 2024
commit cfd30f8610b9e25b003754aac5f6ccf1fbf7bd42
15 changes: 11 additions & 4 deletions torchrl/envs/transforms/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4579,7 +4579,7 @@ def __init__(
self.reset_key = reset_key

# sanity check
for spec in self.primers.values():
for spec in self.primers.values(True, True):
if not isinstance(spec, TensorSpec):
raise ValueError(
"The values of the primers must be a subtype of the TensorSpec class. "
Expand Down Expand Up @@ -4642,9 +4642,16 @@ def transform_observation_spec(
device = observation_spec.device
matteobettini marked this conversation as resolved.
Show resolved Hide resolved
except RuntimeError:
device = self.device
if self.primers.shape[: len(observation_spec.shape)] != observation_spec.shape:
self.primers = self._expand_shape(self.primers)
observation_spec.update(self.primers.to(device))
primers = self.primers.clone()
if self.primers.shape != observation_spec.shape:
try:
# We try to set the primer shape to the observation spec shape
primers.shape = observation_spec.shape
except ValueError:
# If we fail, we expnad them to that shape
primers = self._expand_shape(primers)

observation_spec.update(primers.to(device))
return observation_spec

def transform_input_spec(self, input_spec: TensorSpec) -> TensorSpec:
Expand Down
Loading