💌 Envelope: a JAX-native environment interface

# Create environments from JAX-native suites you have installed, ...
env = envelope.create("gymnax::CartPole-v1")

# ... interact with the environments using a simple interface, ...
state, info = env.init(key)
states, infos = jax.lax.scan(env.step, state, actions)
plt.plot(infos.reward.cumsum())

# ... and enjoy a powerful ecosystem of wrappers.
env = envelope.wrappers.AutoResetWrapper(env)
env = envelope.wrappers.VmapWrapper(env)
env = envelope.wrappers.ObservationNormalizationWrapper(env)

🌍 Simple, expressive interaction!

Environments are pytrees. Squish them through JAX transformations and trace their parameters.
Idiomatic jax-y interface of init(key: Key) -> State, Info and step(state: State, action: PyTree) -> State, Info. You can directly jax.scan over a step(...)!
Spaces are super simple. No Tuple, Dict nonsense! There are two spaces: Continuous and Discrete, which you can compose into a PyTreeSpace.
Explicit episode truncation supports correctly handling bootstrapping for value-function targets.
No auto-reset by default. Resetting every step can be expensive!

💪 Powerful, composable wrappers!

Carry state across episodes to track running statistics, for example to normalize observations.
Explicit wrapper composition keeps episode boundaries correct. See the wrapper compatibility guide for enforced constraints and ordering examples.

🔌 Adapters for existing suites

📦	# 🤖	# 🌍
brax	🕺	12
craftax	🕺	4
gymnax	🕺	24
jumanji	🕺 / 👯	25 / 1
kinetix	🕺	4
mujoco_playground	🕺	54
navix	🕺	41

Total	🕺 / 👯	164 / 1

envelope.create("📦::🌍")

let's you create environments from any of the above!

📝 Testing

Default (no optional adapters deps required): uv run pytest -m "not adapters"
Adapters suite (requires full adapters dependency group):
- uv sync --group adapters
- uv run pytest -m adapters
- Tests for suites unavailable in a partial installation are skipped.

🏗️ Installation

pip install jax-envelope
pip install "jax-envelope[navix]"       # one published adapter
pip install "jax-envelope[adapters]"    # all published adapters

Gymnax and Kinetix are temporarily source-backed development adapters rather than published extras.

💞 Related projects

stoa is a very similar project that provides adapters and wrappers for the jumanji-like interface.
Check out all the great suites we have adapters for! gymnax, brax, jumanji, kinetix, craftax, navix, mujoco_playground.
We will be adding support for jaxmarl and pgx in the future, as soon as we figured out the best ever MARL interface for JAX!

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jax_envelope-0.5.0.tar.gz (104.1 kB view details)

Uploaded Jul 16, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

jax_envelope-0.5.0-py3-none-any.whl (48.0 kB view details)

Uploaded Jul 16, 2026 Python 3

File details

Details for the file jax_envelope-0.5.0.tar.gz.

File metadata

Download URL: jax_envelope-0.5.0.tar.gz
Upload date: Jul 16, 2026
Size: 104.1 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for jax_envelope-0.5.0.tar.gz
Algorithm	Hash digest
SHA256	`7f7ff41e61b69f840f12590033d4948e39c7e477b30bc5c175a5dc42d350a95e`
MD5	`8c7c3013850a251668444339015a5a30`
BLAKE2b-256	`866d8b3ae7cd7aff8b73fb9a167d880e2ceb3e5a28787787055459c81f744e03`

See more details on using hashes here.

Provenance

The following attestation bundles were made for jax_envelope-0.5.0.tar.gz:

Publisher: publish.yml on keraJLi/envelope

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: jax_envelope-0.5.0.tar.gz
- Subject digest: 7f7ff41e61b69f840f12590033d4948e39c7e477b30bc5c175a5dc42d350a95e
- Sigstore transparency entry: 2187026100
- Sigstore integration time: Jul 16, 2026
Source repository:
- Permalink: keraJLi/envelope@fcb7880c1cf07919bcddc8dfe39f8864961c3dd8
- Branch / Tag: refs/tags/v0.5.0
- Owner: https://github.com/keraJLi
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yml@fcb7880c1cf07919bcddc8dfe39f8864961c3dd8
- Trigger Event: release