Add test for GAE + rename RolloutBuffer.dones for clarification

[araffin]

Description

• Added a test for GAE computation
• Add comment for return computation
• Renamed dones for clarity

Motivation and Context

• I have raised an issue to propose this change (required for new features and bug fixes)

see #105

Types of changes

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

Checklist:

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (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.
• I have reformatted the code using make format (required)
• I have checked the codestyle using make check-codestyle and make lint (required)
• I have ensured make pytest and make type both pass. (required)
• I have checked that the documentation builds using make doc (required)

Note: You can run most of the checks using make commit-checks.

Note: we are using a maximum length of 127 characters per line

[araffin]

Note: after some quick trials on CartPole-v1, the formulation self.returns = self.advantages + self.values seems to give me consistently better results.

Need to be confirmed with some Breakout/Pong run I think (or any harder env, maybe HalfCheetahBullet would do the trick)

EDIT: this is kind of confirmed by the error in CI "AssertionError: Mean reward below threshold: 87.60 < 90.00"

[Miffyli]

Could you provide references on what would be the "classical" way of computing returns you mention? While yes, this is not what you usually see, the one used inside GAE computation should be theoretically valid (e.g. if you set gae_lambda to one far end, you get the normal, discounted return).

[araffin]

Could you provide references on what would be the "classical" way of computing returns you mention?

• GAE paper: equation 28
  https://arxiv.org/abs/1506.02438

• Sutton Barto book: "9.3. STOCHASTIC-GRADIENT AND SEMI-GRADIENT METHODS": both MC and TD(0) estimates are shown

I guess using the GAE(lambda) + Value ends up being TD(lambda) target, so it should be fine but need a comment, no?
(but still a bit weird...)

[araffin]

Looking at the original commit (openai/baselines@da99706), the comment says "Compute target value using TD(lambda) estimator, and advantage with GAE(lambda)", so yes it is using TD(lambda) and not MC for value estimation.

The line of code:
https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/ppo1/pposgd_simple.py#L66

[Miffyli]

@araffin

That is very good to know, and also good to know your TD(0) experiment (as I understood it) had lower results (which was to be expected, as I understood TD(lambda) should work better) 👍 . A comment on this would make things clear, and maybe even clarify that this is used because it was used by openai baselines and SB2.

Edit: I mixed TD(0) and TD(1) up, see below.

[araffin]

That is very good to know, and also good to know your TD(0) experiment (as I understood it)

I tested MC estimates (as it is the most common way usually)

that this is used because it was used by openai baselines and SB2.

this is in the docstring but not so clear. I will update it ;)

[araffin]

So, I took some time and found the answer in David Silver lecture 4, slide 47 (page 51 in the PDF) "Telescoping in TD(λ)"), it comes from an non-intuitive telescope: https://www.davidsilver.uk/wp-content/uploads/2020/03/MC-TD.pdf

He shows that $G(\lambda) - V(s)$ leads to the GAE formula.

That is very good to know, and also good to know your TD(0) experiment (as I understood it) had lower results (which was to be expected, as I understood TD(lambda) should work better

TD(1) is MC estimate (non-intuitive when looking at the definition)
TD(0) is one-step estimate with bootstrapping

Edit: from John Schulman tutorial slide 19 (http://joschu.net/docs/2016-NIPS-Tutorial.pdf):

[Miffyli]

Added some clarifications to docs. Two major comments:

1. The GAE-test is to ensure GAE-computation works out correctly, right? Two worries here:
   a) What is the implementation in test based on (i.e. can we trust it?)
   b) I tried changing the code to compute GAE over three episodes (one rollout lasts for max_steps * 3). Only first two tests pass, but remaining do not. I figure this should pass if the implementations were correct?
2. Should dones in compute_returns_and_advantages arguments be renamed to last_episode_starts? This is to clarify dones and self.episode_starts have same meaning, but sadly this mean updating code around the repository to reflect this change.

Sorry for nitpicking. This has been such a headache I want to get it absolutely right this time so it can be left in peace for now :)

[araffin]

a) What is the implementation in test based on (i.e. can we trust it?)

Different things. As you noticed, I took the idea from #105, so that if there is an off-by-one error, the test will fail (I did some testing trying to re-introduce the bug and it failed as expected).

Then, I did some unit testing (GAE(1) is MC estimate) while writing the test (+ manual debugging).
For the lambda-return implementation, I based it on SB3... mostly because we did a lot of testing with that code (and if it is wrong, well, it is wrong but it works ^^").

b) I tried changing the code to compute GAE over three episodes (one rollout lasts for max_steps * 3). Only first two tests pass, but remaining do not. I figure this should pass if the implementations were correct?

Good point. I think I did not write the test for more than one episode, and in fact, it should probably fail without modification (that was to keep things simple and easy to monitor).

hould dones in compute_returns_and_advantages arguments be renamed to last_episode_starts?

I did not rename it because I think it has a different role compared to the episode start (is the very last step of the rollout terminal?), and it is used only once (I think I wrote that in the docstring).

[Miffyli]

Good point. I think I did not write the test for more than one episode, and in fact, it should probably fail without modification (that was to keep things simple and easy to monitor).

Since main code should support this, I think the test should also check for this. In addition, I think it should also test other reward models than "last step rewards 1.0" (e.g. "first step rewards 1.0" and "all steps reward 1.0"), to cover edge-cases and most common cases. I have framework for this and multi-episode test (which fail), which I can push if you'd like.

I did not rename it because I think it has a different role compared to the episode start (is the very last step of the rollout terminal?), and it is used only once (I think I wrote that in the docstring).

Ah right, now I see it 👍 . You are right they have different meanings. Lets leave it as is.

[araffin]

I have framework for this and multi-episode test (which fail), which I can push if you'd like.

I would prefer if you push tests that pass :p

[Miffyli]

Sorry for taking this long! I have added a test for multi-episode rollout that passes. If it checks out, feel free to accept and merge 👍