[Feature] Lightning integration example #2057

svnv-svsv-jm · 2024-04-04T14:24:05Z

Description

This PR offers a convenient lightning.pytorch.LightningModule base class, from which one can inherit to be able to train a torchrl model using lightning.

Motivation and Context

This PR is inspired by this issue: Lightning-Universe/lightning-bolts#986

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds core functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Documentation (update in the documentation)
Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

I have read the CONTRIBUTION guide (required)
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

pytorch-bot · 2024-04-04T14:24:08Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2057

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-04-04T14:24:11Z

Hi @svnv-svsv-jm!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

facebook-github-bot · 2024-04-04T16:06:18Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

vmoens

Pretty impressive work! It'll take me some time to review it.
Cc'ing people who may be interested
@giadefa @albertbou92 @BY571 @tchaton

Approved by mistake

vmoens

I'm leaning towards accepting a lightning backend for the trainers but given the limited bandwidth on my time I won't have much time to scale the PPO trainer to others.

A recipe for other trainers would probably help others code their own trainer. For instance a tutorial (in a second PR)?
Are you considering adding other trainers?

Since this is a new "trainer" I think it should be moved to torchrl/trainers and we will need to make the 2 APIs somewhat compatible.

An example in torchrl/examples would be welcome!

I didn't do a very "in-depth" review, just pieces that I could easily spot for a more homogeneous formatting within the lib!

vmoens · 2024-04-07T11:00:11Z

torchrl/lightning/_base.py

+__all__ = ["BaseRL"]
+
+import typing as ty
+from loguru import logger


unless this is packed with lightning, we won't be using it.
torchrl has a logger under torchrl._utils

Yep, forgot to remove loguru, which is used in the project I copied most of the code from. I will check out torchrl's logger or remove logging entirely for this.

vmoens · 2024-04-07T11:00:36Z

torchrl/lightning/_base.py

@@ -0,0 +1,205 @@
+"""Creates a helper class for more complex models."""
+
+__all__ = ["BaseRL"]


we don't use __all__ but import relevant classes in __init__.py

vmoens · 2024-04-07T11:01:18Z

torchrl/lightning/_base.py

@@ -0,0 +1,205 @@
+"""Creates a helper class for more complex models."""


Missing headers

vmoens · 2024-04-07T11:01:38Z

torchrl/lightning/accelerators.py

@@ -0,0 +1,33 @@
+__all__ = ["find_device"]


missing headers
we don't use __all__

torchrl/lightning/accelerators.py

vmoens · 2024-04-07T15:30:15Z

torchrl/lightning/loops.py

+        this will never be called by the `pl.Trainer` in the `on_train_epoch_end` hook.
+        We have to call it manually in the `training_step`."""
+        scheduler = self.lr_schedulers()
+        assert isinstance(scheduler, LRScheduler)


no assert in the codebase

vmoens · 2024-04-07T15:38:11Z

torchrl/lightning/loops.py

+        batch_idx: int = 0,
+        tag: str = "train",
+    ) -> Tensor:
+        """Common step."""


Can we expand this a bit?

vmoens · 2024-04-07T15:38:29Z

torchrl/lightning/loops.py

+                loss = loss + value
+                loss_dict[f"{key}/{tag}"] = value
+        # Sanity check and return
+        assert isinstance(loss, torch.Tensor)


no assert in codebase

vmoens · 2024-04-07T15:38:55Z

torchrl/lightning/ppo.py

+"""Template for a PPO model on the pendulum env, for Lightning."""
+
+__all__ = ["PPOPendulum"]


missing header, no __all__

vmoens · 2024-04-07T15:39:13Z

torchrl/lightning/ppo.py

+import typing as ty
+
+import torch
+from tensordict import TensorDict  # type: ignore


why the

# type: ignore

?

My mypy in my VSCode was complaining... So I mypy-ignored the line. This is not necessary in the codebase.

Co-authored-by: Vincent Moens <[email protected]>

svnv-svsv-jm · 2024-04-08T09:13:58Z

I will take care of the comments, and add an example under torchrl/examples, but I will leave moving the code to torchrl/trainers for later.

Co-authored-by: Vincent Moens <[email protected]>

vmoens · 2024-06-12T09:02:07Z

@svnv-svsv-jm I think that just moving it to trainers without change of API would already be a good thing!

svnv-svsv-jm · 2024-06-17T21:02:07Z

@svnv-svsv-jm I think that just moving it to trainers without change of API would already be a good thing!

I moved it by allowing from torchrl.trainers import RLTrainingLoop, even though the actual source code is still under torchrl.lightning.

vmoens · 2024-06-19T13:05:43Z

Not sure that cuts it unfortunately. The long term goal is that there is no torchrl.lightning so if we merge this as-if deprecating it will be be hard on the user side

svnv-svsv-jm · 2024-06-20T08:29:42Z

@vmoens No problem! Just wanted to make sure I understood correctly. So indeed I will just move torch.lightning to torchrl.trainers entirely.

passed tests

fdeeb9a

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 Apr 4, 2024

vmoens previously approved these changes Apr 4, 2024

View reviewed changes

vmoens changed the title ~~Lightning integration example~~ [Feature] Lightning integration example Apr 7, 2024

vmoens added the enhancement New feature or request label Apr 7, 2024

vmoens reviewed Apr 7, 2024

View reviewed changes

Use pipe symbol instead of Union

36211ff

Co-authored-by: Vincent Moens <[email protected]>

svnv-svsv-jm and others added 5 commits April 8, 2024 18:40

Use | instead of Union

7fdddca

Co-authored-by: Vincent Moens <[email protected]>

Use | instead of Union

3db2aef

Co-authored-by: Vincent Moens <[email protected]>

Use | instead of Union

48f040f

Co-authored-by: Vincent Moens <[email protected]>

handle pr comments

6f59d8d

add examples + test for it

7775755

allow "from torchrl.trainers import RLTrainingLoop"

394571a

remove "lightning" module and move all to "trainers"

3b767ce

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature] Lightning integration example #2057

[Feature] Lightning integration example #2057

svnv-svsv-jm commented Apr 4, 2024

pytorch-bot bot commented Apr 4, 2024 •

edited

Loading

facebook-github-bot commented Apr 4, 2024

facebook-github-bot commented Apr 4, 2024

vmoens left a comment

vmoens left a comment

vmoens Apr 7, 2024

svnv-svsv-jm Apr 8, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

vmoens Apr 7, 2024

svnv-svsv-jm Apr 8, 2024

svnv-svsv-jm commented Apr 8, 2024

vmoens commented Jun 12, 2024

svnv-svsv-jm commented Jun 17, 2024

vmoens commented Jun 19, 2024

svnv-svsv-jm commented Jun 20, 2024

		@@ -0,0 +1,205 @@
		"""Creates a helper class for more complex models."""

		__all__ = ["BaseRL"]

		"""Template for a PPO model on the pendulum env, for Lightning."""

		__all__ = ["PPOPendulum"]

[Feature] Lightning integration example #2057

Are you sure you want to change the base?

[Feature] Lightning integration example #2057

Conversation

svnv-svsv-jm commented Apr 4, 2024

Description

Motivation and Context

Types of changes

Checklist

pytorch-bot bot commented Apr 4, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2057

facebook-github-bot commented Apr 4, 2024

Action Required

Process

facebook-github-bot commented Apr 4, 2024

vmoens left a comment

Choose a reason for hiding this comment

vmoens left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

svnv-svsv-jm commented Apr 8, 2024

vmoens commented Jun 12, 2024

svnv-svsv-jm commented Jun 17, 2024

vmoens commented Jun 19, 2024

svnv-svsv-jm commented Jun 20, 2024

pytorch-bot bot commented Apr 4, 2024 •

edited

Loading