[BugFix] Fix shifted value computation with an LSTM #2941

vmoens · 2025-05-07T13:11:56Z

Stack from ghstack (oldest at bottom):

-> [BugFix] Fix shifted value computation with an LSTM #2941

[ghstack-poisoned]

pytorch-bot · 2025-05-07T13:12:00Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[PREEMPTIVE] Removal of ephemeral variants on scale-config.yml

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

ghstack-source-id: b912177 Pull-Request-resolved: #2941

[ghstack-poisoned]

ghstack-source-id: ac120cd Pull-Request-resolved: #2941

[ghstack-poisoned]

ghstack-source-id: a9d9bc9 Pull-Request-resolved: #2941

[ghstack-poisoned]

ghstack-source-id: 08b4697 Pull-Request-resolved: #2941

louisfaury · 2025-05-13T17:58:59Z

torchrl/objectives/utils.py

+    # Test:
+    # import torch
+    # from tensordict import TensorDict
+    # a = torch.randn(3, 4)
+    # b = TensorDict(a=torch.randn(3, 4), batch_size=3)
+    # ds = (0, 0)
+    # v0, v1, v2 = zip(*tuple(tree_map(lambda d, x: x.unbind(d), (0, 0), (a, b))))


Should this be removed before merging?

yep i'll make a test out of it

louisfaury · 2025-05-13T18:02:44Z

torchrl/objectives/value/advantages.py

            for i, name in enumerate(data.names):
                if name == "time":
                    ndim = i + 1
                    break


Nit: Should we instead rely on ndim = data.names.index("time") + 1 with a try/catch instead? Would be cleaner imo.

Fair point. The reason I wasn't not doing it is that the try/except isn't compatible with compile but in practice, once we reach the warning we're in a non-compilable code anyway.

louisfaury · 2025-05-13T18:04:20Z

torchrl/objectives/value/advantages.py

@@ -435,62 +441,70 @@ def _call_value_nets(
            value_net = self.value_network
        in_keys = value_net.in_keys
        if single_call:
+            # Reshape to -1 because we cannot guarantee that all dims have the same number of done states


I am not sure I follow why this comment is here.

should be moved down

louisfaury · 2025-05-13T18:18:18Z

torchrl/objectives/value/advantages.py

-                    slice(data.shape[ndim - 1], None),
+            data_copy = data.copy()
+            done = data_copy["next", "done"].clone()
+            done[(slice(None),) * (ndim - 1) + (-1,)].fill_(True)


Suggested change

done[(slice(None),) * (ndim - 1) + (-1,)].fill_(True)

# Mark the last step of every sequence as done.

done[(slice(None),) * (ndim - 1) + (-1,)].fill_(True)

louisfaury · 2025-05-13T18:43:30Z

torchrl/objectives/value/advantages.py

-                idx_ = (slice(None),) * (ndim - 1) + (
-                    slice(data.shape[ndim - 1], None),
+            data_copy = data.copy()
+            done = data_copy["next", "done"].clone()


I'll admit having quite a hard time following the logic that lays below, but more importantly why it is needed. I think a few comments could really help the next reader :)

thomasbbrunner · 2025-05-14T05:47:50Z

test/test_cost.py

+        for p in recurrent_module.parameters():
+            p.data *= 1 + torch.randn_like(p.data) / 10


Curious as to why overwriting the RNN parameters?

I did that to make sure we had enough power to detect slight changes in behavior (ie that every item that could be numerically different than the same computed with the opposite shifted param would be), but in practice it's less useful now

thomasbbrunner · 2025-05-14T05:48:22Z

test/test_cost.py

+            recurrent_module,
+            Mod(mlp_policy, in_keys=["intermediate"], out_keys=["action"]),
+        )
+        # value_net.select_out_keys("state_value")


Suggested change

# value_net.select_out_keys("state_value")

thomasbbrunner · 2025-05-14T05:49:12Z

test/test_cost.py

+        # value_net.select_out_keys("state_value")
+        env = env.append_transform(recurrent_module.make_tensordict_primer())
+        vals = env.rollout(1000, policy_net, break_when_any_done=False)
+        value_net(vals.copy())


I think this call to the value net is not necessary?

Suggested change

value_net(vals.copy())

it's there for init purposes

thomasbbrunner · 2025-05-14T05:51:38Z

test/test_cost.py

+        # TODO: where to put this?
+        vals["next", "is_init"] = vals["is_init"]


Should we solve this TODO before merging?

yes! I think this is something we can safely ask GAE to do for us

thomasbbrunner · 2025-05-14T05:57:41Z

torchrl/modules/tensordict_module/rnn.py

+        _hidden0_in = hidden0_in[..., 0, :, :]
+        _hidden1_in = hidden1_in[..., 0, :, :]


This change is interesting. I'm guessing the shape of hiddenN_in is (batch, sequence, num_lstm_layers, hidden_size).

Were we before not using the first hidden state of a sequence, but the first hidden layer in the LSTM?

it's more an aesthetic change than anything. Just after we swap dims based on negative indices, so using negative indices all the way through helps clarify things. In both cases, we're indexing the 0th element of the sequence.

[ghstack-poisoned]

ghstack-source-id: 9ccbf82 Pull-Request-resolved: #2941

ghstack-source-id: 9ccbf82 Pull-Request-resolved: #2941 (cherry picked from commit 1813e8e)

Update

c5f4ace

[ghstack-poisoned]

vmoens pushed a commit that referenced this pull request May 7, 2025

[BugFix] Fix shifted value computation with an LSTM

b6867dd

ghstack-source-id: b912177 Pull-Request-resolved: #2941

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 May 7, 2025

Update

5cfd491

[ghstack-poisoned]

vmoens pushed a commit that referenced this pull request May 7, 2025

[BugFix] Fix shifted value computation with an LSTM

180fc22

ghstack-source-id: ac120cd Pull-Request-resolved: #2941

Update

2e0872c

[ghstack-poisoned]

vmoens pushed a commit that referenced this pull request May 13, 2025

[BugFix] Fix shifted value computation with an LSTM

d692103

ghstack-source-id: a9d9bc9 Pull-Request-resolved: #2941

vmoens added the bug Something isn't working label May 13, 2025

Update

36c9b2a

[ghstack-poisoned]

vmoens pushed a commit that referenced this pull request May 13, 2025

[BugFix] Fix shifted value computation with an LSTM

3bb38be

ghstack-source-id: 08b4697 Pull-Request-resolved: #2941

louisfaury reviewed May 13, 2025

View reviewed changes

thomasbbrunner reviewed May 14, 2025

View reviewed changes

Update

9f5aaef

[ghstack-poisoned]

vmoens pushed a commit that referenced this pull request May 14, 2025

[BugFix] Fix shifted value computation with an LSTM

4f2b5bf

ghstack-source-id: 9ccbf82 Pull-Request-resolved: #2941

vmoens merged commit 9f5aaef into gh/vmoens/142/base May 14, 2025
50 of 69 checks passed

vmoens pushed a commit that referenced this pull request May 14, 2025

[BugFix] Fix shifted value computation with an LSTM

1813e8e

ghstack-source-id: 9ccbf82 Pull-Request-resolved: #2941

vmoens deleted the gh/vmoens/142/head branch May 14, 2025 08:51

vmoens pushed a commit that referenced this pull request May 16, 2025

[BugFix] Fix shifted value computation with an LSTM

203dfff

ghstack-source-id: 9ccbf82 Pull-Request-resolved: #2941 (cherry picked from commit 1813e8e)

vmoens mentioned this pull request May 20, 2025

GAE does not support LSTM-based value network. #2444

Closed

1 task

	done[(slice(None),) * (ndim - 1) + (-1,)].fill_(True)
	# Mark the last step of every sequence as done.
	done[(slice(None),) * (ndim - 1) + (-1,)].fill_(True)

		for p in recurrent_module.parameters():
		p.data *= 1 + torch.randn_like(p.data) / 10

		# TODO: where to put this?
		vals["next", "is_init"] = vals["is_init"]

		_hidden0_in = hidden0_in[..., 0, :, :]
		_hidden1_in = hidden1_in[..., 0, :, :]

[BugFix] Fix shifted value computation with an LSTM #2941

[BugFix] Fix shifted value computation with an LSTM #2941

Conversation

vmoens commented May 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

vmoens commented May 7, 2025 •

edited

Loading

pytorch-bot bot commented May 7, 2025 •

edited

Loading