search

coax-dev / coax

Modular framework for Reinforcement Learning in python
https://coax.readthedocs.io
MIT License
168 stars 17 forks source link

Fix NStepEntropyRegularizer #12

Closed frederikschubert closed 2 years ago

frederikschubert commented 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.

frederikschubert commented 2 years ago

It also enables the use of other observation_space types like Dict.