v0.1.10

New Features

• Added BatchedEnv ABC class — acts as the contract for any environment that produces N independent agent results per step. VecEnv is one example being the base implementation that uses it. Any downstream packages can ship their own (e.g., composite multi-agent scenes) and slot them into MultiVecEnv without the need for envrax changes.
• Added name property to JaxEnv (defaults to subclass name); Wrapper which delegates to inner environment. Required by MultiEnv / MultiVecEnv for automatic creation of dictionary keys.
• MultiEnv and MultiVecEnv now accept either a list (keys auto-derived from env.name with _0/_1 suffixes on duplicates) or a custom dictionary with personalised keys.

Breaking Changes

• MultiVecEnv is rewritten as a JAX-native, dict-keyed container. State is now Dict[str, chex.ArrayTree] to represent a proper pytree. The cross-env-type (vectorized envs) dispatch loop runs inside a single jax.jit boundary, eliminating per-call Python overhead.
• MultiEnv refactored to match the dict-keyed pattern (dispatch stays Python-loop based — each inner env keeps its own compile cycle).
• make_multi and make_multi_vec now take pre-built env instances (List | Dict) instead of registered names. Construct envs via make / make_vec first, then bundle.
• Removed: reset_at, step_at, class_groups — access inner envs directly via multi.envs[key].
• Renamed: num_envs → n_envs, total_envs → total_slots, vec_envs → envs (on MultiVecEnv).
• VecEnv.render(state, index=0) is now render_slot(state, slot_idx=0); added slot_state(state, slot_idx) for raw single-slot state extraction.

Improvements/Adjustments

• VecEnv auto-reset is replaced with jax.lax.cond and a jnp.where mask. Same semantics, half the HLO size for faster compilation.
• VecEnv.compile() now warms both the typical and done=True reset branches, so the first real episode termination doesn't pay a fresh compile cost.
• Documentation updated to match the new API surface.

v0.1.9

Bug Fixes

• Removed JAX XLA cache config overrides to prevent downstream package issues

v0.1.8

Improvements/Adjustments

• Updated chex.Array -> jax.Array for API consistency and greater IDE support.
• Bumped jax package version from 0.9.2 -> 0.10.0

v0.1.7

Tweaks

• Default XLA complication cache directory now set to <cwd>/.jax_cache for greater project flexibility and monitoring. Path can be overridden by setting ENVRAX_CACHE_DIR

Bug Fixes

• Fixed make_vec and make_multi_vec silently bypassing the persistent cache when pre_warm=False. Cache is now configured whenever jit_compile=True, regardless of pre_warm

v0.1.6

New Features

• Space now declares shape and dtype as part of its contract. Custom spaces with computed shapes use the field(init=False) + __post_init__ pattern (see the custom spaces tutorial).
• New properties on MultiEnv and MultiVecEnv: observation_shapes / action_shapes, observation_sizes / action_sizes, observation_dtypes / action_dtypes (plus single_* variants on MultiVecEnv), and a pad_dims() method returning (max_action_size, max_observation_size) across the fleet.

v0.1.5

Bug Fixes

• Fixed VecEnv.compile() producing dummy actions that only worked for Discrete action spaces.

v0.1.4 — First Usable Release

v0.1.4 — First Usable Release

This release marks the first version of Envrax that's practically usable end-to-end as an API standard. Earlier releases established the shape of the package, but several rough edges in the public surface have been cleaned up so that you can actually pick it up, build your environments and use them without issues.

What's New

Documentation

We've introduced a comprehensive documentation suite to guide you through using the package and getting the most out of it. It includes:

• An Essentials track - this walks you through the key parts of the standard. From understanding the key building blocks, through to building your first environment, vectorizing it, and more!
• An Advanced track for more complicated use cases and customization options
• A "Getting Started" guide with uv / pip / poetry install paths

API Reform

Several breaking changes were made to fix design issues that would have been painful to live with longer-term. Here's some key highlights:

• JaxEnv is now generic with four type parameters. Subclasses now declare JaxEnv[ObsSpaceT, ActSpaceT, StateT, ConfigT] (e.g. BallEnv[Box, Discrete, BallState, BallConfig]), giving full IDE autocomplete and static type-checking on env.observation_space, env.action_space, env.config, and the EnvState returned by reset / step. Previously these returned untyped values, forcing horrible hacky overrides or # type: ignore requirements. As engineers, we are huge fans of clean, easily readable and maintainable code - hacky overrides be damned!
• make() and make_vec() no longer return tuples. Previously they returned (env, config) even though env.config already exposed the resolved config. Now they just return the environment, matching make_multi() / make_multi_vec() for a consistent factory surface and aligning more with the Gymnasium API structure.
• make_multi() / make_multi_vec() lost their config parameter. The old shared config parameter was either silently broken or misleading on heterogeneous fleets. Different environments have different config types so trying to pass a single config across varying environment doesn't make sense!
• They're now true "suite builders" in that they use each environment's registered default config. Per-env overrides are still possible but now handled through manual MultiEnv([make(...), ...]) composition or pre-registered variants.
• RecordVideo now fails fast under JAX transforms. Previously, using it inside jax.jit / jax.vmap / jax.lax.scan would silently produce empty MP4s or confusing trace-time errors. It now detects tracers and raises a clear RuntimeError when first called to catch this issue faster.

Migration

The biggest difference from the previous release is the environment creation:

# Before
class BallEnv(JaxEnv):
    ...

env, config = envrax.make("BallEnv-v0")
vec_env, config = envrax.make_vec("BallEnv-v0", n_envs=64)


# After
class BallEnv(JaxEnv[Box, Discrete, BallState, BallConfig]):
    ...

env = envrax.make("BallEnv-v0")          # config is on env.config
vec_env = envrax.make_vec("BallEnv-v0", n_envs=64)

For per-env config in multi factories, register the variants ahead of time or build MultiEnv / MultiVecEnv manually with make().

Further Development

With this release we are satisfied that the API standard is in a good state without the intention for any further major updates (for now at least). We built Envrax so that we could use it for our own experiments and environment suites, so we'll be doing our own extensive development to really get a feel for it and help weed out any "clunkiness".

We've tried to keep it as lightweight and flexible as possible so you, as the developer, have full control over building your own JAX-native environments that you can share to the world and use for your own research.

If you experience any bugs or want to discuss some potential additions/changes, feel free to open a ticket. Feedback is always welcome!

v0.1.3

Bumped jax version from 0.9.0.1 -> 0.9.2.

v0.1.2

Added JitWrapper, cache compiling and multiple make() factory methods for greater flexibility across envrax suite.

v0.1.1

Merge pull request #3 from Achronus/v0.1.1

V0.1.1