Package: jit-env
jit_env: A Jax Compatible RL Environment APIjit_envis a library that aims to adhere closely to thedm_envinterface
while allowing forjaxtransformations inside Environment implementations
and defining clear type annotations.Likedm_envour API consists of the main components:jit_env.Environment: An abstract base class for RL environments.jit_env.TimeStep: A container class representing the outputs of the environment on each time step (transition).jit_env.specs: A module containing primitives that are used to describe the format of the actions consumed by an environment, as well as the observations, rewards, and discounts it returns.This is extended with the components:jit_env.Wrapper: An interface built on top of Environment that allows modular transformations of the base Environment.jit_env.Action, jit_env.Observation, jit_env.State: Explicit types that concern Agent-Environment IO.jit_env.compat: A Module containing API hooks to other Agent-Environment interfaces likedm_envorgymnasium.jit_env.wrappers: A Module containing a few generally useful implementations forWrapper(that simultaneously serves as a reference).Note that this module is only an interface and does not implement any
Environments itself. The implementations inexamplesserve to illustrate the syntax.
For a more thorough review of the semantics, please refer to thedm-envwiki and compare our implementation ofjit_env.Environmentwithdm_env.Environmentand the conversion as given incompat.py.Installationjit_envcan be installed with (it is recommended to installjaxfirst):python -m pip install jit-envYou can also install it directly from our GitHub repository using pip:python -m pip install git+git://github.com/joeryjoery/jit_env.gitor alternatively by checking out a local copy of our repository and running:python -m pip install /path/to/local/jit_env/The Big Difference withdm_envThe main difference between this API and the standarddm_envAPI is
that our definition ofjit_env.Environmentis functionally pure.
This allows the the logic to e.g., be batched over or accelerated
usingjax.vmaporjax.jit.On top of that, we extend thespecslogic of whatdm_envprovides.
Thespecsmodule defines primitive for how the Agent interacts with
the Environment. We explicitly implement additional specs that are
compatible withjaxbasedPyTreeobjects.
This allows for tree-based operations on thespecobjects themselves,
which in turn gives some added flexibility in designing desired
state-action spaces.Some other modified behaviours include:restartproviding a reference value for reward and discount in place ofNoneStepTypeis no longer anenumtype asjax.jitwould type convertenumtypes to jax primitives anyway. It remains a namespace for defining episode boundaries.TimeStepis now a frozenchex.dataclassto allow usage ofreplacewithin the public API (which is private forNamedTuple).TimeStepcarries an additionalextrasfield to carry optional data (metrics) not shown to the agent.all helperrestart,transition, etc., now take ashapevalue to generate the referencerewardordiscountfields.Whyjit_envI developed this module to have a reliable Environment backend that is less subject
to refactoring changes as other libraries while providing free compatibility to bothjaxtransforms as well as any other popular type of Agent-Environment interface.The hope is that this library will not require much maintenance/ alterations
(aside from some type-hint updates) after an official 1.0.0 release.Cite usIf you are a particularly nice person and this work was useful to you, you can
cite this repository as:@misc{jit_env_2023,author={Joery A. de Vries},title={jit\_env: A Jax interface for reinforcement learning environments},year={2023},url={http://github.com/joeryjoery/jit_env}}ReferencesThis library was heavily inspired by the following libraries:dm-env:https://github.com/deepmind/dm_envjumanji:https://github.com/instadeepai/jumanjigymnax:https://github.com/RobertTLange/gymnaxgymnasium:https://github.com/Farama-Foundation/Gymnasium