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] Max Value Writer #1622

Merged
merged 22 commits into from
Oct 18, 2023
Merged

[Feature] Max Value Writer #1622

Changes from 1 commit
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
a8b29ac
max value writer
albertbou92 Oct 10, 2023
c18db48
max value writer
albertbou92 Oct 10, 2023
3d6d9eb
max value writer
albertbou92 Oct 10, 2023
6f07af1
default key
albertbou92 Oct 10, 2023
a02c2d8
added comments
albertbou92 Oct 11, 2023
2b12160
added comments
albertbou92 Oct 11, 2023
bd32b3c
fix
albertbou92 Oct 11, 2023
b6e1be8
add to docs
albertbou92 Oct 11, 2023
e90f4f0
batched extension
albertbou92 Oct 17, 2023
b383615
test
albertbou92 Oct 17, 2023
2aa659b
test
albertbou92 Oct 17, 2023
ee93f20
test
albertbou92 Oct 17, 2023
23e1888
fix
albertbou92 Oct 17, 2023
466ee57
comments feedback
albertbou92 Oct 18, 2023
7f6ca5f
comments feedback
albertbou92 Oct 18, 2023
8ec741b
comments feedback
albertbou92 Oct 18, 2023
36063dd
comments feedback
albertbou92 Oct 18, 2023
da8eb3a
comments feedback
albertbou92 Oct 18, 2023
846170c
comments feedback
albertbou92 Oct 18, 2023
faad875
Merge branch 'main' into max_val_writer
albertbou92 Oct 18, 2023
48a7f31
Update torchrl/data/replay_buffers/writers.py
vmoens Oct 18, 2023
9d00f10
Update torchrl/data/replay_buffers/writers.py
vmoens Oct 18, 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
fix
  • Loading branch information
@albertbou92
albertbou92 committed Oct 11, 2023
commit bd32b3c4232d261ef46e7d0f6580ac214ac9223d
2 changes: 1 addition & 1 deletion torchrl/data/replay_buffers/writers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ def __init__(self, rank_key=None, **kw) -> None:
def add(self, data: Any) -> int:

ret = None
rank_data = data.get("_data", self._rank_key)
rank_data = data.get("_data").get(self._rank_key)
albertbou92 marked this conversation as resolved.
Show resolved Hide resolved

# Sum the rank key, in case it is a whole trajectory
rank_data = rank_data.sum().item()
Copy link
Contributor

Choose a reason for hiding this comment

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

Is this safe?
Maybe we should document what are the expected shapes for this class, eg

[B, T]

but not

[B1, B2, T]

Another option is to check the number of dimensions of the ranking key OR the name of the last dim of the input tensordict (which should be "time").

Not raising any exception and just doing a plain sum could lead to surprising results I think

Copy link
Contributor Author

@albertbou92 albertbou92 Oct 18, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

I added the first option. Since the ranking value has to be a single float we only allow data of the shape [] and [T] for the add method and [B] and [B, T] for the extend method. If data has a time dimension, we sum along it. If too many dimensions are provided, an error is raised.

I did not go for checking the dimension names because it seemed to restrictive. I don't think time dimension is always labelled

Copy link
Contributor

Choose a reason for hiding this comment

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

Not always but mostly
if you get your data from env.rollout or collector, it will.
If from there you store the data in a rb, it will keep the tag.
But if you reshape or do other stuff it could go away.

Expand Down