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] Fixes to RenameTransform #2442

Merged
merged 4 commits into from
Sep 30, 2024
Merged

[BugFix] Fixes to RenameTransform #2442

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b0afcc2
Fixes to RenameTransform
thomasbbrunner Sep 19, 2024
4622937
amend PR
vmoens Sep 19, 2024
3f620a8
Rewording of docstring
thomasbbrunner Sep 19, 2024
48f1606
fix lint
vmoens Sep 20, 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
Next Next commit
Fixes to RenameTransform
  • Loading branch information
@thomasbbrunner
thomasbbrunner committed Sep 19, 2024
commit b0afcc21e1af33c7ccb2f5177b281764d46cffcd
26 changes: 13 additions & 13 deletions torchrl/envs/transforms/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6713,7 +6713,7 @@ def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
out = tensordict.select(*self.in_keys, strict=not self._missing_tolerance)
for in_key, out_key in zip(self.in_keys, self.out_keys):
try:
tensordict.rename_key_(in_key, out_key)
out.rename_key_(in_key, out_key)
except KeyError:
if not self._missing_tolerance:
raise
Expand Down Expand Up @@ -6743,7 +6743,7 @@ def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
)
for in_key, out_key in zip(self.in_keys_inv, self.out_keys_inv):
try:
out.rename_key_(out_key, in_key)
out.rename_key_(in_key, out_key)
except KeyError:
if not self._missing_tolerance:
raise
Expand All @@ -6752,7 +6752,7 @@ def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
else:
for in_key, out_key in zip(self.in_keys_inv, self.out_keys_inv):
try:
tensordict.rename_key_(out_key, in_key)
tensordict.rename_key_(in_key, out_key)
except KeyError:
if not self._missing_tolerance:
raise
Expand Down Expand Up @@ -6802,29 +6802,29 @@ def transform_output_spec(self, output_spec: Composite) -> Composite:

def transform_input_spec(self, input_spec: Composite) -> Composite:
for action_key in self.parent.action_keys:
if action_key in self.in_keys:
for i, out_key in enumerate(self.out_keys): # noqa: B007
if self.in_keys[i] == action_key:
if action_key in self.out_keys_inv:
for i, in_key_inv in enumerate(self.in_keys_inv): # noqa: B007
if self.out_keys_inv[i] == action_key:
break
else:
# unreachable
raise RuntimeError
input_spec["full_action_spec"][out_key] = input_spec[
input_spec["full_action_spec"][in_key_inv] = input_spec[
"full_action_spec"
][action_key].clone()
if not self.create_copy:
del input_spec["full_action_spec"][action_key]
for state_key in self.parent.full_state_spec.keys(True):
if state_key in self.in_keys:
for i, out_key in enumerate(self.out_keys): # noqa: B007
if self.in_keys[i] == state_key:
if state_key in self.out_keys_inv:
for i, in_key_inv in enumerate(self.in_keys_inv): # noqa: B007
if self.out_keys_inv[i] == state_key:
break
else:
# unreachable
raise RuntimeError
input_spec["full_state_spec"][out_key] = input_spec["full_state_spec"][
state_key
].clone()
input_spec["full_state_spec"][in_key_inv] = input_spec[
"full_state_spec"
][state_key].clone()
if not self.create_copy:
del input_spec["full_state_spec"][state_key]
return input_spec
Expand Down
Loading