search

google / evojax

Apache License 2.0
826 stars 78 forks source link

Issue with BatchNorm layer #60

Open minesh1291 opened 1 year ago

minesh1291 commented 1 year ago

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 commented 1 year ago

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.