[Feature] Implement ensemble reduce #1360
Open
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
vmoens
reviewed
Jul 15, 2023
| @@ -16,17 +16,27 @@ | |||
| UnboundedContinuousTensorSpec, | |||
| ) | |||
| from torchrl.envs.utils import set_exploration_type, step_mdp | |||
| from torchrl.modules import LSTMModule, NormalParamWrapper, SafeModule, TanhNormal | |||
| from torchrl.modules import ( | |||
| EnsembleModule, | |||
| return tensordict_reduced | ||
|
|
||
|
|
||
| class EnsembleModule(TensorDictModuleBase): |
| self, | ||
| in_keys: list[str], | ||
| out_keys: list[str], | ||
| reduce_function: Callable[[torch.Tensor], torch.Tensor] = lambda x: x.min(dim=0).values, |
There was a problem hiding this comment.
The best way to approach this is
Suggested change
| reduce_function: Callable[[torch.Tensor], torch.Tensor] = lambda x: x.min(dim=0).values, | |
| reduce_function: Callable[[torch.Tensor], torch.Tensor] = None, |
then
if reduce_function is None:
reduce_function = lambda x: x.min(dim=0).valuesBut do we really want to have "min" as a default?
There was a problem hiding this comment.
What about Enum to select which op?
REDUCE_OP.min
REDUCE_OP.max
REDUCE_OP.sum
REDUCE_OP.any
REDUCE_OP.all
And then we can ask the user along which dim the op should be done.
That dim would also be the dim along which we do the tensordict indexing.
Comment on lines
+52
to
+54
| assert ( | ||
| len(in_keys) == len(out_keys) == 1 | ||
| ), "Reduce only supports one input and one output" |
There was a problem hiding this comment.
This should be a ValueError
|
@smorad Friendly ping about this :) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This implements a reduction module for use with ensembles.
Motivation and Context
This relies on #1359 to make progress on #1344.
Types 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!