Implemented CrossQ

[danielpalen]

This PR implements CrossQ (https://openreview.net/pdf?id=PczQtTsTIX), a novel off-policy deep RL algorithm that carefully uses batch normalisation and removes target networks to achieve state-of-the-art sample efficiency at a much lower computational complexity, as it does not require large update-to-data-ratios.

Description

This implementation is a PyTorch implementation based on the original JAX implementation (https://github.com/adityab/CrossQ).
The following plot shows that the performance matches the performance reported in the original paper, as well as the performance of the open source SBX implementation provided by the authors (evaluated on 10 seeds).

Context

• I have raised an issue to propose this change (required)
  closes [Feature Request] Implement CrossQ #238

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

[danielpalen]

@araffin in my initial PR it seams like one code style check was failing, sorry about that. I fixed it and it passes on my machine now. I hope it will go through now :)

[araffin]

Could you add at least the multi input policy? (so we can try it in combination with HER)
Only the feature extractor should be changed normally.

And what do you think about adding CnnPolicy?

[danielpalen]

This is a good point. I looked into it and have not yet added it. If I am not mistaken this would also require some changes to the CrossQ train() function. Since, now concatenating and splitting the batches would also require some control flow based on the used policy.
For simplicity sake (for now) and since I did not have time to try and evaluate the multi input policy I did not add that yet.

[araffin]

What do you think about updating create_mlp to allow to pass normalization layer/dropout?

Similar to what is done in DLR-RM/stable-baselines3#1036 and proposed in DLR-RM/stable-baselines3#1069

[danielpalen]

I think this would make sense, because the way I implemented it right now is really that nice.

[araffin]

is that needed? self.actor.set_training_mode(False) is already set above?
or you meant self.actor.set_training_mode(True)?

[danielpalen]

I added more mode calls than needed. The reason was, that I wanted to be very specific which more needs to be used where. I think using the wrong BN mode is one of the big gotchas and sources of error when implementing CrossQ. Since this here should be a PyTorch reference to aid others when they want to implement it by themselves I think it is helpful to make the mode very specific to clear up possible confusion.

[araffin]

Thanks a lot for the implementation =)

I'll try later in the week, but how is it in term of runtime? (SAC vs CrossQ in PyTorch)

[danielpalen]

No worries :)

I just pushed most things you requested. I'll add some more specific responses directly to the questions above.

how is it in term of runtime? (SAC vs CrossQ in PyTorch)

It seems to be quite a but slower than the SAC baseline (and the JAX implementation as well).
for 4M steps, SAC HumanoidStandup took around 12 hours whereas CrossQ took 22 hours. Not sure if there are some PyTorch implementation details that could help with speed.

[araffin]

I'm suspecting something is wrong with the current implementation (I'm currently investigating if it is my changes or not).
My setting:

BipedalWalker-v3:
  n_timesteps: !!float 2e5
  policy: 'MlpPolicy'
  buffer_size: 300000
  gamma: 0.98
  learning_starts: 10000
  policy_kwargs: "dict(net_arch=dict(pi=[256, 256], qf=[1024, 1024]))"

With the RL Zoo cli for both SBX and SB3 (see SBX readme to have support)

python train.py --algo crossq --env BipedalWalker-v3 -P --verbose 0 -param n_envs:30 gradient_steps:30 -n 200000

I'm getting much better results with SBX...
I hope it is not the Adam parameters.

[danielpalen]

Did you figure out what the issue is? I was at ICRA until last week so I didn't have time but if you didn't find it yet I can also have a look.

Before I pushed my last commit I benchmarked it and there the results looked as expected.

[araffin]

Did you figure out what the issue is? I was at ICRA until last week so I didn't have time but if you didn't find it yet I can also have a look.

not yet, I was on holidays...

Before I pushed my last commit I benchmarked it and there the results looked as expected.

I mostly observed the discrepancy on the provided env BipedalWalker-v3 and it seems to be there before my changes.
For the other, I didn't have time yet to launch full benchmark.

One difference currently is the optimizer implementation/arguments, I hope it is the one responsible for it.

[araffin]

nevermind, I did some more systematic tests and I couldn't see any significant difference, the implementation looks good =)

Report: https://wandb.ai/openrlbenchmark/sb3-contrib/reports/SB3-Contrib-CrossQ--Vmlldzo4NTE2MTEx

[danielpalen]

Awesome, let me know if you need anything else :)

[araffin]

let me know if you need anything else :)

sure, I need to find some time to go over it and maybe polish things here and there.
I will probably delay the custom layer integrated in SB3 to later.

[araffin]

minor but still a difference with original implementation: the batch norm stats for the actor are updated at every training step?

(in other word, the delay parameter is not taken into account?)

[danielpalen]

good catch, that's true. I guess the easiest fix for this would be to set self.actor.set_training_mode(update_actor_and_temperature).
What do you think?

[araffin]

i reorganized things and now it should be as in the original implementation

[araffin]

I simplified the network creation (need DLR-RM/stable-baselines3#1975 to be merged with master), added the updated beta for Adam (it had an impact on my small experiments with Pendulum) and fixed a wrong default value for BN momentum that I introduced.