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] COMPOSITE_LP_AGGREGATE env variable #1190

Merged
merged 1 commit into from
Jan 21, 2025
+11 −3
Merged

[Feature] COMPOSITE_LP_AGGREGATE env variable #1190

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 11 additions & 3 deletions tensordict/nn/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,7 +450,13 @@ def _generate_next_value_(name, start, count, last_values):
return name.lower()


_composite_lp_aggregate = _ContextManager()
_composite_lp_aggregate = _ContextManager(
default=(
strtobool(os.getenv("COMPOSITE_LP_AGGREGATE"))
if os.getenv("COMPOSITE_LP_AGGREGATE") is not None
else None
)
)


def composite_lp_aggregate(nowarn: bool = False) -> bool | None:
Expand All @@ -467,9 +473,9 @@ def composite_lp_aggregate(nowarn: bool = False) -> bool | None:
if not nowarn:
warnings.warn(
"Composite log-prob aggregation wasn't defined explicitly and ``composite_lp_aggregate()`` will "
"currently return ``True``. However, from v0.9, this behaviour will change and ``composite_lp_aggregate`` will "
"currently return ``True``. However, from v0.9, this behavior will change and ``composite_lp_aggregate`` will "
"return ``False``. Please change your code accordingly by specifying the aggregation strategy via "
"`tensordict.nn.set_composite_lp_aggregate`.",
"`tensordict.nn.set_composite_lp_aggregate` or via the `COMPOSITE_LP_AGGREGATE` environment variable.",
category=DeprecationWarning,
)
return True
Expand All @@ -483,6 +489,8 @@ class set_composite_lp_aggregate(_DecoratorContextManager):
will be summed into a single tensor with the shape of the root tensordict. This behaviour is being deprecated in favor of
non-aggregated log-probs, which offer more flexibility and a somewhat more natural API (tensordict samples, tensordict log-probs, tensordict entropies).

The value of composite_lp_aggregate can also be controlled through the `COMPOSITE_LP_AGGREGATE` environment variable.

Example:
>>> _ = torch.manual_seed(0)
>>> from tensordict import TensorDict
Expand Down
Loading