Issue with BatchNorm layer

[minesh1291]

while using policy network with BatchNorm layer, getting following error:

ModifyScopeVariableError: Cannot update variable "mean" in "/bn_init" because collection "batch_stats" is immutable. 
 (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ModifyScopeVariableError)

[lerrytang]

Hi, thanks for the question.
BN seems to have a different implementation from other "normal" layers.
Currently, there is no example in evojax that deals with BN layers, and we may need some time to write one.
On the other hand, if you solve the problem first, we'd love to have your PR.