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] Avoid collision of "step_count" key from transform and collector #868

Merged
merged 13 commits into from
Jan 25, 2023
Merged

[BugFix] Avoid collision of "step_count" key from transform and collector #868

Changes from 1 commit
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
de8ccf3
init
vmoens Jan 25, 2023
4a88f5c
Merge branch 'main' into fix_private_collector
vmoens Jan 25, 2023
4d92807
amend
vmoens Jan 25, 2023
6e11a30
amend
vmoens Jan 25, 2023
31dcfe4
lint
vmoens Jan 25, 2023
8b9ed46
markers
vmoens Jan 25, 2023
572e8c2
amend
vmoens Jan 25, 2023
fe294ff
amend
vmoens Jan 25, 2023
f587398
amend
vmoens Jan 25, 2023
47db8d5
amend
vmoens Jan 25, 2023
36c9185
amend
vmoens Jan 25, 2023
6a041b9
amend
vmoens Jan 25, 2023
882a5f6
amend
vmoens Jan 25, 2023
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
@vmoens
vmoens committed Jan 25, 2023
commit 6e11a30ebb9308501cda56fac9ae466e8525b0e8
11 changes: 5 additions & 6 deletions torchrl/envs/transforms/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

import torch
from tensordict.tensordict import TensorDict, TensorDictBase
from tensordict.utils import expand_as_right
from torch import nn, Tensor
from torchrl.data.tensor_specs import (
BinaryDiscreteTensorSpec,
Expand Down Expand Up @@ -2634,7 +2635,7 @@ def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
if out_key in tensordict.keys():
value = tensordict[out_key]
dtype = value.dtype
tensordict[out_key] = value * (~_reset).to(dtype)
tensordict[out_key] = value.masked_fill(expand_as_right(_reset, value), 0.0)
elif in_key == "reward":
# Since the episode reward is not in the tensordict, we need to allocate it
# with zeros entirely (regardless of the _reset mask)
Expand Down Expand Up @@ -2771,11 +2772,9 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
next_step_count = step_count + 1
tensordict.set("step_count", next_step_count)
if self.max_steps is not None:
tensordict.set(
"done",
tensordict.get("done")
| (next_step_count >= self.max_steps).unsqueeze(-1),
)
done = tensordict.get("done")
done = done | (next_step_count >= self.max_steps).unsqueeze(-1)
tensordict.set("done", done)
return tensordict

def transform_observation_spec(
Expand Down