Skip to content

[Feature] Migrate to tensordict.nn.TensorDictModule #700

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

Jump to bottom
Merged
merged 15 commits into from
Nov 25, 2022
Merged

[Feature] Migrate to tensordict.nn.TensorDictModule #700

merged 15 commits into from
Nov 25, 2022

Conversation

@tcbegley
Copy link
Contributor

@tcbegley tcbegley commented Nov 22, 2022 •
edited
Loading

Description

This PR adopts tensordict.nn.TensorDictModule and its variants and eliminates duplicated code. We don't use TensorDictModule as a drop-in replacement since tensordict library has no concept of TensorSpecs. Instead we inherit and let the base classes do the heavy lifting, while the corresponding classes in TorchRL largely add spec validation.

See supporting changes in pytorch/tensordict#66

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 22, 2022
@tcbegley tcbegley mentioned this pull request Nov 22, 2022
Merged
@tcbegley
Copy link
Contributor Author

tcbegley commented Nov 22, 2022

I believe the tests are failing because they require the changes in pytorch/tensordict#66. Not sure what the best way to test this is? Since afaik currently TorchRL does not import anything from tensordict.nn, perhaps that PR can be merged first and then these tests, which pull directly from tensordict main, should pass?

@vmoens vmoens added enhancement New feature or request Refactoring Refactoring of an existing feature labels Nov 22, 2022
@tcbegley tcbegley mentioned this pull request Nov 24, 2022
Merged
tcbegley and others added 3 commits November 24, 2022 11:21
@tcbegley
Delete tests duplicated from tensordict
905f55c
@vmoens
Merge branch 'main' into use-tensordict-nn
3a377ec 
# Conflicts:
#	torchrl/modules/__init__.py
#	torchrl/modules/functional_modules.py
#	torchrl/modules/tensordict_module/common.py
#	torchrl/modules/tensordict_module/probabilistic.py
#	torchrl/modules/tensordict_module/sequence.py
#	torchrl/modules/utils/__init__.py
#	torchrl/trainers/helpers/collectors.py
#	torchrl/trainers/helpers/trainers.py
@vmoens
minor
aa8ac95
@tcbegley tcbegley mentioned this pull request Nov 25, 2022
Closed
tcbegley
tcbegley commented Nov 25, 2022
View reviewed changes
torchrl/modules/tensordict_module/common.py Outdated


class TensorDictModule(nn.Module):
class TensorDictModule(_TensorDictModule):
Copy link
Contributor Author

@tcbegley tcbegley Nov 25, 2022

Choose a reason for hiding this comment

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

I was wondering whether we should consider changing the name here, as it could be confusing if both tensordict and TorchRL have TensorDictModule classes.

I don't have a good name, but I'm thinking something like TensorDictModuleWithSpec to convey the difference?

Copy link
Collaborator

@vmoens vmoens Nov 25, 2022

Choose a reason for hiding this comment

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

Yes that's something I thought about too
Maybe we can drop the TensorDict part (it's implicitly the case since torchrl is all-in into adopting tensordict)
ModuleWithSpec? SpecModule? SafeModule? RLModule?

Copy link
Contributor Author

@tcbegley tcbegley Nov 25, 2022

Choose a reason for hiding this comment

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

Dropping TensorDict part sounds good to me.

Of your suggestions I'm most drawn to SafeModule and SpecModule in that order I think?

Copy link
Collaborator

@vmoens vmoens Nov 25, 2022

Choose a reason for hiding this comment

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

Sorry for the slow answer
We can also think about user experience, and for that RLModule seems good
In ReAgent they have a (improperly named) ModelBase class that all models inherit from. It's not only bad because Model-Based in a subfield in RL, but also because it fails to account for what it is (the base class of all policies and value networks).
SafeModule is also good as the main difference is mainly that putting a spec allows for safe sampling / mapping (it's more about safety than about the spec IMO). In that sense RLModule is less self-explanatory.
TLDR let's go for SafeModule :p

@codecov
Copy link

codecov bot commented Nov 25, 2022 •
edited
Loading

Codecov Report

Merging #700 (a814bdf) into main (d27d070) will increase coverage by 0.04%.
The diff coverage is 99.20%.

@@            Coverage Diff             @@
##             main     #700      +/-   ##
==========================================
+ Coverage   88.77%   88.81%   +0.04%     
==========================================
  Files         122      121       -1     
  Lines       21151    20490     -661     
==========================================
- Hits        18776    18198     -578     
+ Misses       2375     2292      -83
Flag Coverage Δ
habitat-gpu 24.81% <82.89%> (+0.44%) ⬆️
linux-cpu 84.82% <97.60%> (-0.06%) ⬇️
linux-gpu 85.65% <98.80%> (-0.09%) ⬇️
linux-jumanji 30.05% <82.89%> (+0.74%) ⬆️
linux-outdeps-gpu 71.62% <85.25%> (-0.46%) ⬇️
linux-stable-cpu 84.67% <97.60%> (-0.04%) ⬇️
linux-stable-gpu 85.40% <98.00%> (-0.01%) ⬇️
linux_examples-gpu 42.98% <75.67%> (+0.24%) ⬆️
macos-cpu 84.48% <97.60%> (-0.07%) ⬇️
olddeps-gpu 74.84% <79.68%> (+0.09%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files Coverage Δ
torchrl/modules/__init__.py 100.00% <ø> (ø)
torchrl/modules/models/exploration.py 91.79% <ø> (ø)
torchrl/modules/planners/cem.py 97.61% <ø> (ø)
test/test_tensordictmodules.py 97.79% <97.75%> (-0.10%) ⬇️
test/smoke_test.py 100.00% <100.00%> (ø)
test/test_collector.py 98.60% <100.00%> (ø)
test/test_cost.py 96.95% <100.00%> (ø)
test/test_env.py 98.73% <100.00%> (ø)
test/test_exploration.py 98.84% <100.00%> (ø)
test/test_functorch.py 94.59% <100.00%> (ø)
... and 32 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

vmoens
vmoens approved these changes Nov 25, 2022
View reviewed changes
Copy link
Collaborator

@vmoens vmoens left a comment

Choose a reason for hiding this comment

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

Great stuff! A lot less code: the best code is no code at all

vmoens
vmoens approved these changes Nov 25, 2022
View reviewed changes
Copy link
Collaborator

@vmoens vmoens left a comment

Choose a reason for hiding this comment

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

A couple of suggestions for the docstrings
It's ok that things are copied for the rest i guess

@tcbegley tcbegley mentioned this pull request Nov 25, 2022
Merged
vmoens
vmoens approved these changes Nov 25, 2022
View reviewed changes
Copy link
Collaborator

@vmoens vmoens left a comment

Choose a reason for hiding this comment

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

Cool!

@vmoens vmoens merged commit a4506df into pytorch:main Nov 25, 2022
@tcbegley tcbegley deleted the use-tensordict-nn branch March 29, 2023 17:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vmoens vmoens vmoens approved these changes

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. enhancement New feature or request Refactoring Refactoring of an existing feature

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@tcbegley @vmoens @facebook-github-bot