Algorithm with recurrent policy

[eleninisioti]

I was wondering how you would extend an existing algorithm, let's say es to allow for a recurrent neural network instead of a feedforward one. You would need to change make_policy_network to replace the MLP with a newly created RNN so my impression is that you'd need to change only networks.py and not train.py. But not sure how to implement it

[btaba]

Besides the make_policy_network, I'd expect you have to modify the actor_step to handle the recurrent state

https://github.com/google/brax/blob/16304037a36b1d9c8c0b3084f57d1159627b636b/brax/training/acting.py#L34-L55