# `jit_env`: A Jax Compatible RL Environment API
`jit_env` is a library that aims to adhere closely to the `dm_env` interface
while allowing for `jax` transformations inside Environment implementations
and defining clear type annotations.
Like `dm_env` our 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 like `dm_env` or `gymnasium`.
- `jit_env.wrappers`: A Module containing a few generally useful implementations for `Wrapper` (that simultaneously serves as a reference).
Note that this module is only an interface and does not implement any
Environments itself. The implementations in `examples` serve to illustrate the syntax.
For a more thorough review of the semantics, please refer to the [dm-env](https://github.com/deepmind/dm_env/blob/master/dm_env/specs.py)
wiki and compare our implementation of `jit_env.Environment` with `dm_env.Environment` and the conversion as given in [`compat.py`](https://github.com/joeryjoery/jit_env/blob/main/jit_env/compat.py).
## Installation
`jit_env` can be installed with (it is recommended to install `jax` first):
`python -m pip install jit-env`
You can also install it directly from our GitHub repository using pip:
`python -m pip install git+git://github.com/joeryjoery/jit_env.git`
or alternatively by checking out a local copy of our repository and running:
`python -m pip install /path/to/local/jit_env/`
## The Big Difference with `dm_env`
The main difference between this API and the standard `dm_env` API is
that our definition of `jit_env.Environment` is functionally pure.
This allows the the logic to e.g., be batched over or accelerated
using `jax.vmap` or `jax.jit`.
On top of that, we extend the `specs` logic of what `dm_env` provides.
The `specs` module defines primitive for how the Agent interacts with
the Environment. We explicitly implement additional specs that are
compatible with `jax` based `PyTree` objects.
This allows for tree-based operations on the `spec` objects themselves,
which in turn gives some added flexibility in designing desired
state-action spaces.
Some other modified behaviours include:
- `restart` providing a reference value for reward and discount in place of `None`
- `StepType` is no longer an `enum` type as `jax.jit` would type convert `enum` types to jax primitives anyway. It remains a namespace for defining episode boundaries.
- `TimeStep` is now a frozen `chex.dataclass` to allow usage of `replace` within the public API (which is private for `NamedTuple`).
- `TimeStep` carries an additional `extras` field to carry optional data (metrics) not shown to the agent.
- all helper `restart`, `transition`, etc., now take a `shape` value to generate the reference `reward` or `discount` fields.
### Why `jit_env`
I developed this module to have a reliable Environment backend that is less subject
to refactoring changes as other libraries while providing free compatibility to both `jax`
transforms 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 us
If you are a particularly nice person and this work was useful to you, you can
cite this repository as:
```bibtex
@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}
}
```
## References
This library was heavily inspired by the following libraries:
- dm-env: [https://github.com/deepmind/dm_env](https://github.com/deepmind/dm_env)
- jumanji: [https://github.com/instadeepai/jumanji](https://github.com/instadeepai/jumanji)
- gymnax: [https://github.com/RobertTLange/gymnax](https://github.com/RobertTLange/gymnax)
- gymnasium: [https://github.com/Farama-Foundation/Gymnasium](https://github.com/Farama-Foundation/Gymnasium)
Raw data
{
"_id": null,
"home_page": "https://github.com/joeryjoery/jit_env",
"name": "jit-env",
"maintainer": "",
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": "",
"keywords": "reinforcement-learning python machine learning jax",
"author": "Joery A. de Vries",
"author_email": "J.A.deVries@tudelft.nl",
"download_url": "https://files.pythonhosted.org/packages/13/59/019f0454dfa0c18695b3f1240d11f8c49953de30907fda4ae0a38ebeb162/jit_env-0.1.4.tar.gz",
"platform": null,
"description": "\n\n\n# `jit_env`: A Jax Compatible RL Environment API\n`jit_env` is a library that aims to adhere closely to the `dm_env` interface\nwhile allowing for `jax` transformations inside Environment implementations\nand defining clear type annotations.\n\nLike `dm_env` our API consists of the main components:\n\n- `jit_env.Environment`: An abstract base class for RL environments.\n- `jit_env.TimeStep`: A container class representing the outputs of the environment on each time step (transition).\n- `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.\n\nThis is extended with the components:\n- `jit_env.Wrapper`: An interface built on top of Environment that allows modular transformations of the base Environment.\n- `jit_env.Action, jit_env.Observation, jit_env.State`: Explicit types that concern Agent-Environment IO.\n- `jit_env.compat`: A Module containing API hooks to other Agent-Environment interfaces like `dm_env` or `gymnasium`.\n- `jit_env.wrappers`: A Module containing a few generally useful implementations for `Wrapper` (that simultaneously serves as a reference).\n\nNote that this module is only an interface and does not implement any\nEnvironments itself. The implementations in `examples` serve to illustrate the syntax.\nFor a more thorough review of the semantics, please refer to the [dm-env](https://github.com/deepmind/dm_env/blob/master/dm_env/specs.py) \nwiki and compare our implementation of `jit_env.Environment` with `dm_env.Environment` and the conversion as given in [`compat.py`](https://github.com/joeryjoery/jit_env/blob/main/jit_env/compat.py).\n\n## Installation\n`jit_env` can be installed with (it is recommended to install `jax` first):\n\n`python -m pip install jit-env`\n\nYou can also install it directly from our GitHub repository using pip:\n\n`python -m pip install git+git://github.com/joeryjoery/jit_env.git`\n\nor alternatively by checking out a local copy of our repository and running:\n\n`python -m pip install /path/to/local/jit_env/`\n\n## The Big Difference with `dm_env`\nThe main difference between this API and the standard `dm_env` API is\nthat our definition of `jit_env.Environment` is functionally pure.\nThis allows the the logic to e.g., be batched over or accelerated \nusing `jax.vmap` or `jax.jit`. \n\nOn top of that, we extend the `specs` logic of what `dm_env` provides.\nThe `specs` module defines primitive for how the Agent interacts with \nthe Environment. We explicitly implement additional specs that are \ncompatible with `jax` based `PyTree` objects.\nThis allows for tree-based operations on the `spec` objects themselves, \nwhich in turn gives some added flexibility in designing desired \nstate-action spaces.\n\nSome other modified behaviours include: \n - `restart` providing a reference value for reward and discount in place of `None`\n - `StepType` is no longer an `enum` type as `jax.jit` would type convert `enum` types to jax primitives anyway. It remains a namespace for defining episode boundaries.\n - `TimeStep` is now a frozen `chex.dataclass` to allow usage of `replace` within the public API (which is private for `NamedTuple`).\n - `TimeStep` carries an additional `extras` field to carry optional data (metrics) not shown to the agent.\n - all helper `restart`, `transition`, etc., now take a `shape` value to generate the reference `reward` or `discount` fields.\n\n### Why `jit_env`\nI developed this module to have a reliable Environment backend that is less subject\nto refactoring changes as other libraries while providing free compatibility to both `jax` \ntransforms as well as any other popular type of Agent-Environment interface. \n\nThe hope is that this library will not require much maintenance/ alterations \n(aside from some type-hint updates) after an official 1.0.0 release. \n\n## Cite us\nIf you are a particularly nice person and this work was useful to you, you can\ncite this repository as:\n```bibtex\n@misc{jit_env_2023,\n author={Joery A. de Vries},\n title={jit\\_env: A Jax interface for reinforcement learning environments},\n year={2023},\n url={http://github.com/joeryjoery/jit_env}\n}\n```\n\n## References\nThis library was heavily inspired by the following libraries:\n- dm-env: [https://github.com/deepmind/dm_env](https://github.com/deepmind/dm_env)\n- jumanji: [https://github.com/instadeepai/jumanji](https://github.com/instadeepai/jumanji)\n- gymnax: [https://github.com/RobertTLange/gymnax](https://github.com/RobertTLange/gymnax)\n- gymnasium: [https://github.com/Farama-Foundation/Gymnasium](https://github.com/Farama-Foundation/Gymnasium)\n",
"bugtrack_url": null,
"license": "MIT",
"summary": "A Jax interface for Reinforcement Learning environments.",
"version": "0.1.4",
"project_urls": {
"Homepage": "https://github.com/joeryjoery/jit_env"
},
"split_keywords": [
"reinforcement-learning",
"python",
"machine",
"learning",
"jax"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "041efea019b1309bfcc3f9b51426745243f324b598fb3800a8d1832cc26e59f1",
"md5": "de0c713f82d39d7c077ca1f7db31cc69",
"sha256": "c62ab81eb8a7ba8324914a0435658400bc6fc1505b242452d8aad8644101a55e"
},
"downloads": -1,
"filename": "jit_env-0.1.4-py3-none-any.whl",
"has_sig": false,
"md5_digest": "de0c713f82d39d7c077ca1f7db31cc69",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 33095,
"upload_time": "2023-05-12T13:32:04",
"upload_time_iso_8601": "2023-05-12T13:32:04.275816Z",
"url": "https://files.pythonhosted.org/packages/04/1e/fea019b1309bfcc3f9b51426745243f324b598fb3800a8d1832cc26e59f1/jit_env-0.1.4-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "1359019f0454dfa0c18695b3f1240d11f8c49953de30907fda4ae0a38ebeb162",
"md5": "3e72ea27ae1942a473aa918eec98fa93",
"sha256": "1553c415753f51633e8f3885af988b4d7ea9b3163fd16391bf751749ab04e97e"
},
"downloads": -1,
"filename": "jit_env-0.1.4.tar.gz",
"has_sig": false,
"md5_digest": "3e72ea27ae1942a473aa918eec98fa93",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 31291,
"upload_time": "2023-05-12T13:32:07",
"upload_time_iso_8601": "2023-05-12T13:32:07.836310Z",
"url": "https://files.pythonhosted.org/packages/13/59/019f0454dfa0c18695b3f1240d11f8c49953de30907fda4ae0a38ebeb162/jit_env-0.1.4.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2023-05-12 13:32:07",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "joeryjoery",
"github_project": "jit_env",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [],
"tox": true,
"lcname": "jit-env"
}
JSON
