[Feature Request] Add Attention nets (GTrXL model in particular)

[RemiG3]

🚀 Feature Request

This feature request is a duplicate from stable-baselines3 (see DLR-RM/stable-baselines3#177).

The idea is to add the GTrXL model in the contrib repo from the paper Stabilizing Transformers for Reinforcement Learning, as done in RLlib: https://github.com/ray-project/ray/blob/master/rllib/models/torch/attention_net.py.

@araffin has already mentioned that he created it and will make it public (comment).

I wonder if this is still relevant?

[araffin]

@araffin has already mentioned that he created it and will make it public (DLR-RM/stable-baselines3#177 (comment)).

I meant the SB3 contrib repo.

For GTrXL, are you willing to contribute that algorithm?
Please read carefully the contributing guide if you decide to.

[RemiG3]

I meant the SB3 contrib repo.

Sorry for the misunderstanding.

For GTrXL, are you willing to contribute that algorithm?

I'm not sure yet, I will try to implement it for my experiments first.

[araffin]

Also related: https://github.com/maohangyu/TIT_open_source

[richardjozsa]

@RemiG3 hey, have you started to implement it? Mayba I can give a free hand in it :)

[RemiG3]

Yes, I have implemented it, but not tested properly. I'm currently having some troubles with my custom environment that I'm trying to solve.

@araffin is it possible to create a new branch for this feature (to share the code)?
If it is possible, I'll clean up the code and push it to this new branch soon.

[araffin]

Yes, I have implemented it, but not tested properly. I'm currently having some troubles with my custom environment that I'm trying to solve.

@araffin is it possible to create a new branch for this feature (to share the code)? If it is possible, I'll clean up the code and push it to this new branch soon.

yes, that's what a fork and pull request are meant for

[richardjozsa]

I have came accross on this, this is quite modular and easy to tune, Transformers-RL, the only backside is that, it has been implemented only to gaussian policy.

[RemiG3]

Hey, I finally made the PR #176 to share the code.

It should work, but I'm not sure about the performances.
It would be nice if someone could make comparisons with other methods (or RLlib attention net for example).
I won't have time these next days.

[eric000888]

RemiG3, Thank you for adding attention net to contrib. what's the shape of the input would be look like , for example if I want to use cartpole environment?
Thanks again.

[RemiG3]

Thank you, @eric000888, for reporting this (feel free to provide the code you tested as you did in your first edits).

I have updated the branch to fix a bug on the dimension of minibatchs.
But, I still have an exception when batch_size = 1 or n_steps = 1 and I found the same exception for RecurrentPPO.

So, it should now work for batch_size > 1 and n_steps > 1 (as RecurrentPPO).

EDIT: I also add assertions about these cases, as in the original PPO.