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] patch rand_action in TransformedEnv to read the base_env method #2699

Merged
merged 5 commits into from
Jan 26, 2025
+22 −0
Merged

[BugFix] patch rand_action in TransformedEnv to read the base_env method #2699

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
764e2e0
Update
vmoens Jan 17, 2025
2794f70
Update
vmoens Jan 17, 2025
55934fb
Update
vmoens Jan 21, 2025
1024d61
Update
vmoens Jan 25, 2025
e4076f4
Update
vmoens Jan 26, 2025
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
22 changes: 22 additions & 0 deletions torchrl/envs/transforms/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -795,6 +795,28 @@ def input_spec(self) -> TensorSpec:
input_spec = self.__dict__.get("_input_spec", None)
return input_spec

def rand_action(self, tensordict: Optional[TensorDictBase] = None) -> TensorDict:
if type(self.base_env).rand_action is not EnvBase.rand_action:
# TODO: this will fail if the transform modifies the input.
# For instance, if an env overrides rand_action and we build a
# env = PendulumEnv().append_transform(ActionDiscretizer(num_intervals=4))
# env.rand_action will NOT have a discrete action!
# Getting a discrete action would require coding the inverse transform of an action within
# ActionDiscretizer (ie, float->int, not int->float).
Copy link
Collaborator

@kurtamohler kurtamohler Jan 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Is there a reason we couldn't use self.action_spec.rand()?

>>> import torchrl
>>> env = torchrl.envs.PendulumEnv().append_transform(torchrl.envs.ActionDiscretizer(num_intervals=4))
>>> env.action_spec.rand()
tensor([3])
>>> env.action_spec.rand().dtype
torch.int64

Copy link
Contributor Author

Choose a reason for hiding this comment

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

yes, what I meant with that comment is that if your base env redefines the rand_action then the action you'll get won't be transformed

import torchrl
Pendulum = torchrl.envs.PendulumEnv
rand_action = Pendulum.rand_action
Pendulum.rand_action = lambda *args, **kwargs: rand_action(*args, **kwargs)

env = Pendulum().append_transform(torchrl.envs.ActionDiscretizer(num_intervals=4))

print(env.action_spec.rand())
print(env.action_spec.rand().dtype)
print(env.rand_action())

which prints

tensor([3])
torch.int64
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

# We can loosely check that the action_spec isn't altered - that doesn't mean the action is
# intact but it covers part of these alterations.
#
# The following check may be expensive to run and could be cached.
if self.full_action_spec != self.base_env.full_action_spec:
raise RuntimeError(
f"The rand_action method from the base env {self.base_env.__class__.__name__} "
"has been overwritten, but the transforms appended to the environment modify "
"the action. To call the base env rand_action method, we should then invert the "
"action transform, which is (in general) not doable."
)
return self.base_env.rand_action(tensordict)
return super().rand_action(tensordict)

def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
# No need to clone here because inv does it already
# tensordict = tensordict.clone(False)
Expand Down
Loading