Closed frederikschubert closed 2 years ago
The current implementation of the NStepEntropyRegularizer fails when updating to jax==0.2.26. This PR improves the implementation to use vmap and fixes the problem.
NStepEntropyRegularizer
jax==0.2.26
vmap
It also enables the use of other observation_space types like Dict.
observation_space
Dict
The current implementation of the
NStepEntropyRegularizer
fails when updating tojax==0.2.26
. This PR improves the implementation to usevmap
and fixes the problem.