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

[Feature] Rework to_one_hot and to_categorical to take a tensor as parameter #816

Merged
merged 7 commits into from
Feb 20, 2023
Merged

[Feature] Rework to_one_hot and to_categorical to take a tensor as parameter #816

Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ecefb66
Rework MultOneHotDiscreteTensorSpec to_categorical_method
riiswa Jan 10, 2023
e05f9db
Rework MultiDiscreteTensorSpec to_onehot method
riiswa Jan 13, 2023
e1eec1d
Rework OneHotDiscreteTensorSpec & DiscreteTensorSpec to_onehot & to_c…
riiswa Jan 13, 2023
4e6b5b5
Add docs
riiswa Jan 13, 2023
02a7c8e
Remove to numpy method
riiswa Jan 23, 2023
4b875a6
Fix device in MultiDiscreteTensorSpec.is_in
riiswa Feb 15, 2023
5f3e311
Separe conversion functions for spec and for tensor
riiswa Feb 17, 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
Rework MultiDiscreteTensorSpec to_onehot method
  • Loading branch information
@riiswa
riiswa committed Feb 15, 2023
commit e05f9db98ba54234a055b941ff28ecb02f5aa9a7
13 changes: 10 additions & 3 deletions torchrl/data/tensor_specs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1219,8 +1219,9 @@ def _project(self, val: torch.Tensor) -> torch.Tensor:
vals = self._split(val)
return torch.cat([super()._project(_val) for _val in vals], -1)

def to_categorical(self, val: Optional[torch.Tensor] = None, safe: bool = True) -> Union[
MultiDiscreteTensorSpec, torch.Tensor]:
def to_categorical(
self, val: Optional[torch.Tensor] = None, safe: bool = True
) -> Union[MultiDiscreteTensorSpec, torch.Tensor]:
if val is None:
return MultiDiscreteTensorSpec(
[_space.n for _space in self.space],
Expand Down Expand Up @@ -1510,7 +1511,13 @@ def to_onehot(self, val: Optional[torch.Tensor] = None, safe: bool = True) -> Un
else:
if safe:
self.assert_is_in(val)
return val
return torch.cat(
[
torch.nn.functional.one_hot(val[..., i], n)
for i, n in enumerate(self.nvec)
],
-1,
)

def expand(self, *shape):
if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)):
Expand Down