Recurrent PPO

[araffin]

Description

Experimental version of PPO with LSTM policy.

Current status: usable but not polished, see #53 (comment)

Missing:

• benchmark https://wandb.ai/sb3/no-vel-envs/reports/PPO-vs-RecurrentPPO-aka-PPO-LSTM-on-environments-with-masked-velocity--VmlldzoxOTI4NjE4 (and models in https://wandb.ai/openrlbenchmark/sb3?workspace=user-araffin )
• more tests (reproducibility + identity learning/callback)
• reduce code duplication in buffer
• document buffer + rest
• more testing of the different architecture (shared vs separate lstm for actor/critic)

Known issue: if the model was train on GPU and tested on CPU, a warning will be issued because it cannot unpickle the lstm initial states. This is ok as they will be reset anyway in setup_model() and it doesn't affect prediction.

Context

• I have raised an issue to propose this change (required)
  closes recurrent policy implementation in ppo [feature-request] DLR-RM/stable-baselines3#18

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)
• The functionality/performance matches that of the source (required for new training algorithms or training-related features).
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have included an example of using the feature (required for new features).
• I have included baseline results (required for new training algorithms or training-related features).
• I have updated the documentation accordingly.
• I have updated the changelog accordingly (required).
• 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)

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

[araffin]

@andrewwarrington

When doing the backprop through the action probabilities (here), you just take a single step in the RNN, i.e. the gradient with respect to the RNN parameters given just the previous RNN state and input?

Actually not, see the shape in the buffer https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/blob/feat/ppo-lstm/sb3_contrib/common/recurrent/buffers.py#L164
I just pushed proper masking for the padded sequences. We do backprop for the sequences collected during the n_steps of the rollout.

[HamiltonWang]

for the time being, is there any way we can feedback the older data back to the trainer as input to mimic a crude version of LSTM? any sample code to do that?

[araffin]

for the time being, is there any way we can feedback the older data back to the trainer as input to mimic a crude version of LSTM? any sample code to do that?

Hello,
you can already use this PR if you need to use recurrent PPO (see install from source in our doc), otherwise you can use frame stacking or history wrapper (see code in the RL Zoo).

[henrydeclety]

when LSTM for A2C?

[araffin]

when LSTM for A2C?

a2c is a special case of ppo ;) (cc @vwxyzjn )

[vwxyzjn]

@henrydeclety see https://github.com/vwxyzjn/a2c_is_a_special_case_of_ppo. We have a paper coming out soon...

[vwxyzjn]

The preprint of the paper is out at https://arxiv.org/abs/2205.09123 @henrydeclety :)

[HamiltonWang]

Hello, you can already use this PR if you need to use recurrent PPO (see install from source in our doc), otherwise you can use frame stacking or history wrapper (see code in the RL Zoo).

I’ll give it a try

[EloyAnguiano]

How could I configure the maximum sequence length for the LSTM?

[philippkiesling]

@EloyAnguiano As far as I could tell from the code, the implementation in SB3 does not have a sequence length, but saves the hidden state between steps of your environment and then uses it as input. So the maximum sequence length for the lstm would be the number of steps (n_steps) before you update your policy.

This way you only need to compute each input once, instead of refeeding it every new step.