[Feature] AbsorbingStateTransform #2290
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2290
Note: Links to docs will display an error until the docs builds have been completed. ❌ 5 New Failures, 1 Unrelated FailureAs of commit 1d43d8b with merge base 8e43ac8 ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
Will update the docstring with examples. |
| >>> from torchrl.envs import GymEnv | ||
| >>> t = AbsorbingStateTransform(max_episode_length=1000) | ||
| >>> base_env = GymEnv("HalfCheetah-v4") | ||
| >>> env = TransformedEnv(base_env, t) |
There was a problem hiding this comment.
That's not very informative about the functionality ;)
| terminate_key: Optional[NestedKey] = "terminated", | ||
| ): | ||
| if in_keys is None: | ||
| in_keys = "observation" # default |
| batch_size = observation.size(0) | ||
| if self._done: | ||
| # Create absorbing states for the batched observations | ||
| absorbing_state = torch.eye(observation.size(1) + 1)[-1] |
There was a problem hiding this comment.
This is rather wasteful, we creating a big tensor and indexing it, plus this is a view on a storage hence the original storage isn't cleared when you index.
Besides it lacks dtype and device.
You can create an incomplete eye with m and n, see the doc here
| # Create absorbing states for the batched observations | ||
| absorbing_state = torch.eye(observation.size(1) + 1)[-1] | ||
| return absorbing_state.expand(batch_size, -1) | ||
| zeros = torch.zeros(batch_size, 1) |
There was a problem hiding this comment.
Missing device and dtype
You could use observation.new_zeros
| @@ -8557,3 +8557,157 @@ def _inv_call(self, tensordict): | |||
| if self.sampling == self.SamplingStrategy.RANDOM: | |||
| action = action + self.jitters * torch.rand_like(self.jitters) | |||
| return tensordict.set(self.in_keys_inv[0], action) | |||
|
|
|||
|
|
|||
| class AbsorbingStateTransform(ObservationTransform): | |||
There was a problem hiding this comment.
We need tests for this class
It should be registered in the __init__.py and put in the doc.
| def forward(self, tensordict: TensorDictBase) -> TensorDictBase: | ||
| raise RuntimeError(FORWARD_NOT_IMPLEMENTED.format(type(self))) | ||
|
|
||
| def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
there is a version of this that works for all batch sizes. This one will only work with uni of bidimensional batch sizes.
| elif observation.dim() == 2: | ||
| # Batched observations | ||
| batch_size = observation.size(0) | ||
| if self._done: |
There was a problem hiding this comment.
do we need an in-place value? How does that work if one sub-env is done and the other not?
Maybe we could read the done state and change it on the fly, without using local attribute
| ) | ||
| return tensordict | ||
| done = tensordict.get(self.done_key) | ||
| self._done = done.any() |
There was a problem hiding this comment.
this means that if any sub-env is done all are done?
| # Single observation | ||
| if self._done: | ||
| # Return absorbing state which is [0, ..., 0, 1] | ||
| return torch.eye(observation.size(0) + 1)[-1] |
There was a problem hiding this comment.
what if the observation is more than 1d?
vmoens
added
enhancement
Description
Adds AbsorbingStateTransform as used in the DAC paper.
Motivation and Context
Why is this change required? What problem does it solve?
If it fixes an open issue, please link to the issue here.
You can use the syntax
close #15213if this solves the issue #15213Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!