jit_env

This release fixes a number of stale issues and updates the typing annotations in core.

This includes;

• Addition of ExtendObservation wrapper to nest TimeStep fields inside the observation (useful for POMDPs)
• Update all calls to jax.random.PRNGKey to jax.random.key.
• Some typing updates in core to make use of TypeVars. This is still incomplete and tracked in #28.
• Update all chex tests to remove ignore_nones arguments.

This release adds a Gymnasium converter for both Environment and Spec-Space.

It also fixes a bug inside the Tree spec and unpack_spec function that was incorrectly flattening and unflattening the custom defined PyTrees.

This release adds a StopGradient wrapper for the Environment Step function. This is a generally useful utility for online reinforcement learning algorithms.

This version also fixes not being able to use State, Action, or Observation as TypeVars as they were supposed to. These are now explicility imported from _core.py rather than defining TypeAliases.

This release removes the render method as an abstractmethod, and moves the rendering responsibility to an external argument dependency.

This release allows mypy to type check this library when installed from PyPI.

Changes:

• Added py.typed inside jit_env and setup.py
• Extended repr in Environment and Wrapper and implemented more detailed reprs in jit_env.wrappers.py.
• Added tests to ensure repr and str work as expected.

This release 0.1.0 covers the entire library adequately with tests to ensure reliable behaviour.

Future Development

Some features related to third-party library compatibility may be added later, such as gymnasium, or maybe brax or gymnax. When these additional features are implemented and tested, this library version will be bumped to 1.0.0.

From this point on, features will only be added in the form of Pull Requests to prevent very sudden API changes.

The 0.0.1 Alpha version was just uploaded to PyPI.

The non-alpha version will be released when wrappers.py and compat.py are 100% covered by tests.

Then, when pytest and mypy give the OK, version 0.0.1 will be released.