🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures, 1 Unrelated Failure

As of commit ddfce88 with merge base 2b8450c ():

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

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 91. Improved: $\large\color{#35bf28}9$. Worsened: $\large\color{#d91a1a}9$.

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 94. Improved: $\large\color{#35bf28}7$. Worsened: $\large\color{#d91a1a}2$.

@AechPro I did not change the n_steps argument as you suggested for consistency within torchrl.
I'm open to change both but I'm not sure how to proceed to make things non-bc breaking

I'll ping you when the preview of the doc is built for you to check if things make sense, in the meantime you can look at the diff

@vmoens I understand that backwards compatibility is important, but I think if we expect torchrl to be largely used by a research community it seems a bit strange for the n_steps argument to behave in a slightly unexpected way.

Consider the case where a user of torchrl is re-implementing an algorithm like RAINBOW in good faith, and they naturally set the n_step argument to 3 as suggested by the paper. Using the current behavior this would actually produce a multi-step return estimate with n=4 according to the algorithm implemented in the paper, which would meaningfully change the performance of the algorithm. I can imagine our hypothetical torchrl user becoming quite confused when they cannot replicate the results of RAINBOW despite having (seemingly) implemented everything correctly here.

Further, if we expect new algorithms to be written using torchrl, this sort of discrepancy in the meaning of the multi-step parameter could introduce a point of confusion in the opposite direction: if a person not using torchrl is attempting to implement a paper which describes results from an algorithm implemented with torchrl using the multi-step return object, that person may find it difficult to replicate the results of the torchrl algorithm for the same reason as our hypothetical RAINBOW user.

In my opinion preserving the meaning of the parameter in the multi-step return algorithm is more important than backwards compatibility because of the two examples I presented above, but I understand others may feel differently. At the very least I would strongly advocate for some easy to see documentation about this behavior whenever there is a tutorial using the MultiStep object and in the docstring for the object itself so we can mitigate these two potential issues.

Makes sense! Then let's make this n_step congruent with what's expected by the community

Here's the doc https://docs-preview.pytorch.org/pytorch/rl/2008/reference/generated/torchrl.envs.transforms.rb_transforms.MultiStepTransform.html#torchrl.envs.transforms.rb_transforms.MultiStepTransform

I'm having a little trouble understanding the output of the example. We have the MultiStepTransform with n_steps=3, then we sample some timesteps and slice the first 5 entries from our replay buffer. The first entry in that slice is reporting a step count of 9, which I assume is supposed to be equal to 5+n, which would be 5 + 3 = 8 but that would still be an incorrect behavior. The first timestep emitted by the transform should be timestep 1 (or zero if you start counting at zero), because that contains the state from which we're computing the n-step return, so I would expect rb[:]["step_count"][:, 0] to be 1 (or zero) and the final entry at rb[:]["step_count"][:, -1] to be T-n+1 because the internal buffers should only need to contain n-1 waiting timesteps unless there is a terminal state at the end of the internal buffer, in which case it should just compute all of the remaining possible returns at each of those timesteps.

I'm sure I must be misunderstanding what's happening in the example. Could you clarify?

Let me clarify:
Here we look at the step count at the root:

>>> print("step_count", rb[:]["step_count"][:, :5])
step_count tensor([[[ 9],
         [10],
         [11],
         [12],
         [13]],

        [[12],
         [13],
         [14],
         [15],
         [16]]])

Env 0 has steps [9, 10, 11...] and env 1 [12, 13, 14,...]

Then we look at the "next" entry to see the shift.
Without MultiStep, we would have [10, 11, 12, ...] for env 0 and [13, 14, 15, ...] for env 1
Because we use multi-step with a shift of 3 we have [13, 14, 15, ...] and [16, 17, 18, ...] resp.

>>> print("next step_count", rb[:]["next", "step_count"][:, :5])
next step_count tensor([[[13],
         [14],
         [15],
         [16],
         [17]],

        [[16],
         [17],
         [18],
         [19],
         [20]]])

Note that we're looking at the replay buffer content so it doesn't really matter what those values are (ie it's expected that it doesn't start at 0).

For a single env and n=3, you would have this data structure accessible in the buffer

done:        [0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
step_count:        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
next, step_count_orig:        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
next, step_count:        [4, 5, 6, 7, 8, 9, 10, 10, 10, 10]

If you think this isn't the desired behaviour I'd be happy to read what you think would be expected, but to me it looks like what multi-step requires.

Ah-hah, my apologies for misunderstanding. Thanks for the clarification, this looks good!

Although if you are looking to implement the change to the n_step parameter we spoke about earlier right now then I believe the example with 1 env should be as follows:

done:                        [0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
step_count:                  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
next, step_count_orig:       [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
next, step_count:            [3, 4, 5, 6, 7, 8, 9, 10, 10, 10]

Because starting from state 0 we will compute the return estimate as the sum of rewards starting with timestep 0, then 1, then 2, and finally the learning algorithm will need to bootstrap from the state encountered at timestep 3.

So given what you're saying, what it should be is:
n_steps=0 => error
n_steps=1 => no transform
n_steps=2 => 1 shift in the future
etc

I will change that thx

I'm not sure what the right outcome for n_steps=0 is. It seems reasonable to think we should set all of the rewards to zero and then set all the next states to the current states (i.e. step_count = t and next, step_count = t) because this would turn the equation return_estimates = batch["rewards"] + batch["next]["gammas"] * value_estimator(batch["next"]["observation"]) into return_estimates = 0 + 1 * value_estimator(batch["next"]["observation"]) = value_estimator(batch["observation"]) provided we also set the gamma values to 1. This makes some amount of sense because the value of n_steps can be taken to mean the number of reward terms we need to incorporate before using an estimator of the value function in our return estimate, so zero there would just lead us to use our value estimator completely with no reward terms at all

With that said I'm not aware of anything in the literature using n=0, and I wonder if it would be confusing or if it's even necessary.

The way I see it n=0 is equivalent to "doing nothing" which can be achieved by... doing nothing haha

LOL 😅 yeah it makes sense to do it that way. I was just imagining a scenario where maybe a user is doing some sort of hyper-parameter investigation and they wanted to vary the value of n from 0 to some number without changing anything in the underlying learning algorithm, which is maybe a realistic thing to want if someone is interested in measuring the impact of rewards on the return estimate.

I guess they will have to start from 1 :)

I think accounting for those edge cases causes more hustle and loads the doc while bringing little value in practice (+ requires proper testing of a behaviour that is even poorly defined on our end!)

Yeah, fair enough. Let's keep it as you suggested then!

cc @agarwl this is an implementation of multi-step that allows to dynamically change the horizon during training