RLax Files

What's Changed

• rlax: Replace rlax categorical cross entropy computation with distrax components. by @copybara-service in https://github.com/deepmind/rlax/pull/57
• Bugfix to quantile_expected_sarsa. by @copybara-service in https://github.com/deepmind/rlax/pull/63
• Update Jinja2 versioning to avoid Sphinx failures. by @copybara-service in https://github.com/deepmind/rlax/pull/66
• Add test for squashed gaussian in rlax distributions. by @copybara-service in https://github.com/deepmind/rlax/pull/68
• Update squashed gaussian distribution in rlax for prob and logprob to numerically match distrax's implementation. by @copybara-service in https://github.com/deepmind/rlax/pull/69
• Migrate RLax squashed gaussian to use Distrax. Explicitly broadcast shapes in Distrax scalar affine to avoid rank promotion errors. by @copybara-service in https://github.com/deepmind/rlax/pull/70
• Add a particular pair of transforms used by muzero that combine a non linear squashing function with a reparametrisation of the scalar as linear combination of two hot values in a discrete suppport. by @copybara-service in https://github.com/deepmind/rlax/pull/73
• Support Array lambda_ in Vtrace. by @copybara-service in https://github.com/deepmind/rlax/pull/71
• Send deprecation warning for rlax.distributions in favor of using distrax. by @copybara-service in https://github.com/deepmind/rlax/pull/74
• Send deprecation warning for rlax nested_updates in favor of using optax. by @copybara-service in https://github.com/deepmind/rlax/pull/75
• Move usages of soon to be deprecated rlax.periodic_update to optax.periodic_update. by @copybara-service in https://github.com/deepmind/rlax/pull/77
• Add a pair of transforms where the scalar values are reparametrised as the linear combination of two-hot values on a non-linearly spaced discrete support. by @copybara-service in https://github.com/deepmind/rlax/pull/78
• Add moving averages helpers to rlax. by @copybara-service in https://github.com/deepmind/rlax/pull/79
• Update .pylintrc by @copybara-service in https://github.com/deepmind/rlax/pull/80
• Add utilities to extract overlapping subsequences from trajectories. by @copybara-service in https://github.com/deepmind/rlax/pull/81
• Minor edits to moving averages. by @copybara-service in https://github.com/deepmind/rlax/pull/83
• Add utilities to support interruptions. by @copybara-service in https://github.com/deepmind/rlax/pull/84
• Create new version 0.1.3 of RLax. by @copybara-service in https://github.com/deepmind/rlax/pull/86
• Remove incremental_update from rlax: all usages ported to optax.incremental_update by @copybara-service in https://github.com/deepmind/rlax/pull/85
• Pin numpy version <1.23 until new jax version is released, fixing bug that makes mpo_ops_test fail. by @copybara-service in https://github.com/deepmind/rlax/pull/92
• Fix a bug in tree_split_leaves(): squeeze the right axis in case of keepdim=False. by @copybara-service in https://github.com/deepmind/rlax/pull/94
• Fix max_start_idx argument. by @copybara-service in https://github.com/deepmind/rlax/pull/87
• Release a new rlax verison. by @copybara-service in https://github.com/deepmind/rlax/pull/96

Full Changelog: https://github.com/deepmind/rlax/compare/v0.1.2...v0.1.4