[Feature] adding tensor classes annotation for loss functions #1905

SandishKumarHN · 2024-02-13T00:56:05Z

Description

followup from this pull request
copy past:

We project on using https://github.com/Tensorclass to represent losses.

The advantage of tensorclass for losses instead of tensordict is that it will help us use all the features of tensordict while preserving type annotation or even completion.

Changes:
Check the out_keys of the loss;
Create a tensorclass with the respective fields;
Type the forward as returning that class (and/or a tensordict)
Add an argument to return the class in the constructor with the False value by default;
Update the docstrings (not done)
Write a little test to check that things work as expected (this test should be new and not parametrized - if we add one more parameter to the existing tests the code will be much longer and harder to follow, and the tests will run for a long time).

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-02-13T00:56:09Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 14 New Failures

As of commit 9b5f4e6 with merge base 87f3437 ():

NEW FAILURES - The following jobs have failed:

Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
Workflow failed! Resource not accessible by integration
Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
Workflow failed! Resource not accessible by integration
Habitat Tests on Linux / tests (3.9, 11.6) / linux-job (gh)
RuntimeError: Command docker exec -t b6e60aa69af5dedf42f02bc6f5f9da53457d6e9f3602e70ed5117ff208b4a29a /exec failed with exit code 139
Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-cpu (3.8) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-gpu (3.10, 12.1) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-optdeps (3.10, 12.1) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on MacOS CPU / tests (3.11) / macos-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on MacOS CPU / tests (3.8) / macos-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]
Unit-tests on Windows / unittests-cpu / windows-job (gh)
test/test_cost.py::TestDQN::test_distributional_dqn_reduction[sum-9]

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

vmoens

Looks promising!
I'd suggest adding the args to the docstrings and make a single test for each loss that we can get the tensorclass and the we can access the losses inf return_tensorclass is True

vmoens · 2024-02-13T21:31:13Z

torchrl/objectives/a2c.py

@@ -234,6 +250,7 @@ def __init__(
        functional: bool = True,
        actor: ProbabilisticTensorDictSequential = None,
        critic: ProbabilisticTensorDictSequential = None,
+        return_tensorclass: bool = False,


Should be added to the docstrings

working on it.

@vmoens added doctests for tensorclass changes. but I see some doctest issues and blockers. can you please help me resolve.

there are some existing doctest failures, we might need a separate task to address.

what would be the aggregate_loss for each loss within tensorclass?

there are some existing errors like

```Cannot interpret 'torch.int64' as a data type```

```'key "action_value" not found in TensorDict with keys [\'done\', \'logits\', \'observation\', \'reward\', \'state_value\', \'terminated\']' ```

```NameError: name 'actor' is not defined```

etc

# Conflicts: # torchrl/objectives/ppo.py

Subscribers: Tasks: Tags:

SandishKumarHN · 2024-02-24T05:09:58Z

@vmoens can you review once, build errors on resource not related to the PR.

vmoens

Thanks!

New classes should be manually added to the doc
See docs/source/reference/objectives.rst

The feature seems untested to me, am I right?

vmoens · 2024-02-28T22:09:37Z

test/test_cost.py

+
            loss_function="l2",
+
            delay_value=delay_value,
+


not sure why we need these

vmoens · 2024-02-28T22:11:43Z

torchrl/objectives/a2c.py

@@ -33,6 +35,34 @@
 )


+class LossContainerBase:
+    """ContainerBase class loss tensorclass's."""


That isn't very explicit. We should say what this class is about.

Also I think it should live in the common.py file.

I'm also wondering if we should not just make the base a tensorclass and inherit from it without creating new tensorclasses?

If I try to make the base a tensorclass getting below error.

********************************************************************** File "/home/sandish/rl/torchrl/objectives/a2c.py", line 144, in a2c.A2CLoss Failed example: loss(data) Exception raised: Traceback (most recent call last): File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/doctest.py", line 1334, in __run exec(compile(example.source, filename, "single", File "<doctest a2c.A2CLoss[21]>", line 1, in <module> loss(data) File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1561, in _call_impl result = forward_call(*args, **kwargs) File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/tensordict/_contextlib.py", line 126, in decorate_context return func(*args, **kwargs) File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/tensordict/nn/common.py", line 291, in wrapper return func(_self, tensordict, *args, **kwargs) File "/home/sandish/rl/torchrl/objectives/a2c.py", line 503, in forward return A2CLosses._from_tensordict(td_out) File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/tensordict/tensorclass.py", line 327, in wrapper raise ValueError( ValueError: Keys from the tensordict ({'loss_entropy', 'loss_objective', 'entropy', 'loss_critic'}) must correspond to the class attributes (set()).

vmoens · 2024-02-28T22:12:41Z

torchrl/objectives/a2c.py

+    @property
+    def aggregate_loss(self):
+        return self.loss_critic + self.loss_objective + self.loss_entropy


No need to recode this

vmoens · 2024-02-28T22:13:21Z

torchrl/objectives/a2c.py

+    def aggregate_loss(self):
+        result = 0.0
+        for key in self.__dataclass_attr__:
+            if key.startswith("loss_"):
+                result += getattr(self, key)
+        return result


Should be a property
Should always return a tensor
Something like

result = torch.zeros((), device=self.device) ... return result

Missing docstring for this method.

vmoens · 2024-02-28T22:15:20Z

torchrl/objectives/a2c.py

@@ -455,7 +497,7 @@ def _cached_detach_critic_network_params(self):
        return self.critic_network_params.detach()

    @dispatch()
-    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+    def forward(self, tensordict: TensorDictBase) -> A2CLosses:


Suggested change

def forward(self, tensordict: TensorDictBase) -> A2CLosses:

def forward(self, tensordict: TensorDictBase) -> A2CLosses | TensorDictBase:

vmoens · 2024-02-28T22:26:49Z

torchrl/objectives/sac.py

+class LossContainerBase:
+    """ContainerBase class loss tensorclass's."""
+
+    __getitem__ = TensorDictBase.__getitem__
+
+    def aggregate_loss(self):
+        result = 0.0
+        for key in self.__dataclass_attr__:
+            if key.startswith("loss_"):
+                result += getattr(self, key)
+        return result
+
+
+@tensorclass
+class SACLosses(LossContainerBase):
+    """The tensorclass for The SACLoss Loss class."""
+
+    loss_actor: torch.Tensor
+    loss_value: torch.Tensor
+    loss_qvalue: torch.Tensor
+    alpha: torch.Tensor | None = None
+    loss_alpha: torch.Tensor | None = None
+    entropy: torch.Tensor | None = None


vmoens · 2024-02-28T22:27:03Z

torchrl/objectives/sac.py

@@ -541,7 +581,7 @@ def out_keys(self, values):
        self._out_keys = values

    @dispatch
-    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+    def forward(self, tensordict: TensorDictBase) -> SACLosses:


Suggested change

def forward(self, tensordict: TensorDictBase) -> SACLosses:

def forward(self, tensordict: TensorDictBase) -> SACLosses | TensorDictBase:

vmoens · 2024-02-28T22:27:17Z

torchrl/objectives/sac.py

+            out["loss_value"] = loss_value.mean()
+        return TensorDict(out, [])


Why this change?

vmoens · 2024-02-28T22:27:34Z

torchrl/objectives/td3.py

+class LossContainerBase:
+    """ContainerBase class loss tensorclass's."""
+
+    __getitem__ = TensorDictBase.__getitem__
+
+    def aggregate_loss(self):
+        result = 0.0
+        for key in self.__dataclass_attr__:
+            if key.startswith("loss_"):
+                result += getattr(self, key)
+        return result
+
+
+@tensorclass
+class TD3Losses(LossContainerBase):
+    """The tensorclass for The TD3 Loss class."""
+
+    loss_actor: torch.Tensor
+    loss_qvalue: torch.Tensor
+    target_value: torch.Tensor | None = None
+    state_action_value_actor: torch.Tensor | None = None
+    pred_value: torch.Tensor | None = None
+    next_state_value: torch.Tensor | None = None


vmoens · 2024-02-28T22:27:48Z

torchrl/objectives/td3.py

@@ -453,7 +492,7 @@ def value_loss(self, tensordict):
        return loss_qval, metadata

    @dispatch
-    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+    def forward(self, tensordict: TensorDictBase) -> TD3Losses:


Suggested change

def forward(self, tensordict: TensorDictBase) -> TD3Losses:

def forward(self, tensordict: TensorDictBase) -> TD3Losses | TensorDictBase:

vmoens

Thanks!

New classes should be manually added to the doc
See docs/source/reference/objectives.rst

The feature seems untested to me, am I right?

SandishKumarHN · 2024-02-29T19:48:44Z

@vmoens address most of your comments above, but doctests are failing with below error not caused by this PR changes.

File "/home/sandish/rl/torchrl/objectives/cql.py", line 128, in cql.CQLLoss
Failed example:
    loss = CQLLoss(actor, qvalue)
Exception raised:
    Traceback (most recent call last):
      File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/doctest.py", line 1334, in __run
        exec(compile(example.source, filename, "single",
      File "<doctest cql.CQLLoss[16]>", line 1, in <module>
        loss = CQLLoss(actor, qvalue)
      File "/home/sandish/rl/torchrl/objectives/cql.py", line 321, in __init__
        self.convert_to_functional(
      File "/home/sandish/rl/torchrl/objectives/common.py", line 289, in convert_to_functional
        params.apply(
      File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/tensordict/nn/params.py", line 125, in new_func
        out = meth(*args, **kwargs)
      File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/tensordict/base.py", line 3824, in apply
        return self._apply_nest(
      File "/home/sandish/.conda/envs/torch_rl/lib/python3.9/site-packages/tensordict/_td.py", line 659, in _apply_nest
        out = TensorDict(
    TypeError: __init__() got an unexpected keyword argument 'filter_empty'

  File "/pytorch/rl/env/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1704, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'DDPGLoss' object has no attribute 'reduction'

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>

vmoens

Great progress!
I'd like to brainstorm about naming of those classes, A2CLosses seems a bit awkward when we have A2CLoss which is a totally different thing.
I don't think we should rename "loss" in "loss_objective" as part of this PR

torchrl/objectives/a2c.py

vmoens · 2024-03-12T18:48:58Z

torchrl/objectives/a2c.py

@@ -300,6 +322,8 @@ def __init__(
        if gamma is not None:
            raise TypeError(_GAMMA_LMBDA_DEPREC_ERROR)
        self.loss_critic_type = loss_critic_type
+        self.return_tensorclass = return_tensorclass
+        self.reduction = reduction


torchrl/objectives/common.py

vmoens · 2024-03-12T18:49:57Z

torchrl/objectives/a2c.py

@@ -32,6 +34,15 @@
    VTrace,
 )

+@tensorclass


doens't it work if we make the base class a tensorclass?

Yes, It doesn't work.

vmoens · 2024-03-12T18:51:02Z

torchrl/objectives/common.py

@@ -298,7 +309,7 @@ def _compare_and_expand(param):
        # set the functional module: we need to convert the params to non-differentiable params
        # otherwise they will appear twice in parameters
        with params.apply(
-            self._make_meta_params, device=torch.device("meta"), filter_empty=False
+            self._make_meta_params, device=torch.device("meta")


why is this removed?

torchrl/objectives/common.py

vmoens · 2024-03-12T18:52:45Z

torchrl/objectives/dqn.py

+    loss_objective: torch.Tensor
+    loss: torch.Tensor


torchrl/objectives/dqn.py

vmoens

Great progress!
I'd like to brainstorm about naming of those classes, A2CLosses seems a bit awkward when we have A2CLoss which is a totally different thing.
I don't think we should rename "loss" in "loss_objective" as part of this PR

review changes

SandishKumarHN · 2024-03-14T20:30:55Z

@vmoens made changes based on your review, I still reduction is not being added to the test_cost.py file so all of the failures are related to that.

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 Feb 13, 2024

SandishKumarHN force-pushed the tensorclass-losses branch 3 times, most recently from 3c3bc29 to 0e993e4 Compare February 13, 2024 04:34

adding tensor classes annotation for loss functions

e4761a3

SandishKumarHN force-pushed the tensorclass-losses branch from 0e993e4 to e4761a3 Compare February 13, 2024 05:18

vmoens added the enhancement New feature or request label Feb 13, 2024

vmoens changed the title ~~adding tensor classes annotation for loss functions~~ [Feature] adding tensor classes annotation for loss functions Feb 13, 2024

vmoens reviewed Feb 13, 2024

View reviewed changes

SandishKumarHN added 3 commits February 16, 2024 17:36

review changes add doctests for tensorclass

5d432c8

review changes add doctests for tensorclass | merging to main

387953f

review changes add doctests for tensorclass | merging to main

8e16b63

SandishKumarHN force-pushed the tensorclass-losses branch from 7d6e08f to 8e16b63 Compare February 17, 2024 02:00

vmoens and others added 12 commits February 21, 2024 15:11

Merge remote-tracking branch 'origin/main' into tensorclass-losses

bfb5930

# Conflicts: # torchrl/objectives/ppo.py

amend

bfde82f

amend

7473445

review changes add doctests for tensorclass | merging to main

8163f90

build error fix

b21e43e

Reviewers:

60e3d51

Subscribers: Tasks: Tags:

Merge remote-tracking branch 'upstream/main' into tensorclass-losses

44a70e6

build error fix

23ef8ea

Merge remote-tracking branch 'upstream/main' into tensorclass-losses

191ab2e

build error fix, doc test aggregated func

1e373ca

build error fix, docstring formatted

3f058e1

build error fix, docstring formatted

582c9c5

Merge remote-tracking branch 'upstream/main' into tensorclass-losses

79d8a29

vmoens reviewed Feb 28, 2024

View reviewed changes

flake8 errors

64837f9

SandishKumarHN force-pushed the tensorclass-losses branch from 66fc382 to db00e02 Compare February 29, 2024 19:42

review changes - 1

715d4c0

SandishKumarHN force-pushed the tensorclass-losses branch from db00e02 to 715d4c0 Compare February 29, 2024 19:46

SandishKumarHN and others added 3 commits March 11, 2024 14:49

Merge remote-tracking branch 'upstream/main' into tensorclass-losses

5bb8894

compiler errors

7c0ae77

Update torchrl/objectives/decision_transformer.py

e9125fb

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>

vmoens reviewed Mar 12, 2024

View reviewed changes

compiler errors

bae4237

SandishKumarHN force-pushed the tensorclass-losses branch from c7cf4b9 to bae4237 Compare March 12, 2024 19:00

review changes

e17c91e

review changes

SandishKumarHN force-pushed the tensorclass-losses branch from 95bad71 to e17c91e Compare March 14, 2024 17:37

SandishKumarHN added 2 commits March 14, 2024 14:18

Merge remote-tracking branch 'upstream/main' into tensorclass-losses

f07c4f4

review changes

8b5e0ff

SandishKumarHN force-pushed the tensorclass-losses branch 2 times, most recently from ad2b6c3 to 8b5e0ff Compare March 18, 2024 16:12

SandishKumarHN added 2 commits March 18, 2024 09:15

Merge remote-tracking branch 'upstream/main' into tensorclass-losses

73a4dcd

review changes

9b5f4e6

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature] adding tensor classes annotation for loss functions #1905

[Feature] adding tensor classes annotation for loss functions #1905

SandishKumarHN commented Feb 13, 2024 •

edited

pytorch-bot bot commented Feb 13, 2024 •

edited

vmoens left a comment

vmoens Feb 13, 2024

SandishKumarHN Feb 13, 2024

SandishKumarHN Feb 17, 2024 •

edited

SandishKumarHN commented Feb 24, 2024

vmoens left a comment

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

SandishKumarHN Feb 29, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens Feb 28, 2024

vmoens left a comment

SandishKumarHN commented Feb 29, 2024 •

edited

vmoens left a comment

vmoens Mar 12, 2024

vmoens Mar 12, 2024

SandishKumarHN Mar 14, 2024

vmoens Mar 12, 2024

vmoens Mar 12, 2024

vmoens left a comment

SandishKumarHN commented Mar 14, 2024

	def forward(self, tensordict: TensorDictBase) -> A2CLosses:
	def forward(self, tensordict: TensorDictBase) -> A2CLosses \| TensorDictBase:

		out["loss_value"] = loss_value.mean()
		return TensorDict(out, [])

[Feature] adding tensor classes annotation for loss functions #1905

Are you sure you want to change the base?

[Feature] adding tensor classes annotation for loss functions #1905

Conversation

SandishKumarHN commented Feb 13, 2024 • edited

Description

Types of changes

Checklist

pytorch-bot bot commented Feb 13, 2024 • edited

🔗 Helpful Links

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

❌ 14 New Failures

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

SandishKumarHN Feb 17, 2024 • edited

Choose a reason for hiding this comment

SandishKumarHN commented Feb 24, 2024

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vmoens left a comment

Choose a reason for hiding this comment

SandishKumarHN commented Feb 29, 2024 • edited

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

vmoens left a comment

Choose a reason for hiding this comment

SandishKumarHN commented Mar 14, 2024

SandishKumarHN commented Feb 13, 2024 •

edited

pytorch-bot bot commented Feb 13, 2024 •

edited

SandishKumarHN Feb 17, 2024 •

edited

SandishKumarHN commented Feb 29, 2024 •

edited