Add typed envelope core, wrappers, and adapter support

[keraJLi]

Summary

• Introduce the core envelope package with typed environment, space, struct, and wrapper abstractions.
• Add adapter integrations for Brax, Craftax, Gymnax, Jumanji, Kinetix, MuJoCo Playground, and Navix.
• Expand wrapper support for flattening, normalization, autoreset, vectorization, truncation, statistics, and state injection.
• Refresh docs, release notes, packaging metadata, and CI/publish workflows for the 0.5 release.
• Add broad coverage across adapters, spaces, structs, environment behavior, and wrapper contracts.

Testing

• Added and updated unit and integration tests across adapters, spaces, environment primitives, and wrappers.
• Validation is covered by the expanded test suite and release workflow updates; no separate manual run was requested.

[Copilot]

Pull request overview

Introduces Envelope 0.5’s typed core and expands the wrapper/adapter ecosystem, with a focus on stricter lifecycle semantics (terminal snapshots, wrapper ordering/constraints), improved JAX compatibility, and updated packaging/CI for release distribution.

Changes:

• Adds stack-constraint validation for wrapper composition plus new wrapper semantics around terminal transitions (final / final_valid) and state injection.
• Expands/refreshes adapters (Brax, Craftax, Gymnax, Jumanji, MuJoCo Playground, Navix, Kinetix) with captured default horizons and backend-metadata normalization.
• Updates spaces/struct utilities (canonical Container extras ordering, static-field validation, broadened contains/sampling behavior) and adds extensive tests + release/CI workflows.

Reviewed changes

Copilot reviewed 75 out of 78 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/wrappers/test_vmap_wrapper.py	Updates vmap/normalization composition assertions.
tests/wrappers/test_truncation_wrapper.py	Aligns truncation tests with new “OR-with-inner-truncation” semantics and validation.
tests/wrappers/test_state_injection_wrapper.py	Updates state injection tests to require full reset_info and validates nested wrapper structures.
tests/wrappers/test_stack_constraints.py	New tests for directional stack constraints and marker mixins.
tests/wrappers/test_pooled_init_vmap_wrapper.py	Adds pooled-init semantics/compatibility tests (final placeholders, incompatibilities, metadata).
tests/wrappers/test_observation_normalization_wrapper.py	Expands normalization tests: pytrees, broadcast-aware stats, final normalization, pooled wrapping.
tests/wrappers/test_normalization.py	Updates RMV tests for per-leaf counts and safer std computation.
tests/wrappers/test_hardening_contracts.py	New cross-wrapper lifecycle/contract tests (keywords, jittability, ordering).
tests/wrappers/test_flatten_observation_wrapper.py	Adds test for flattening outside vmap preserving batch prefix.
tests/wrappers/test_flatten_action_wrapper.py	Makes flatten-action step JIT-compatible; adds batch-prefix preservation test.
tests/wrappers/test_episode_statistics_wrapper.py	Updates episode stats reset semantics and vmap composition expectations.
tests/wrappers/test_environment_wrapper.py	Adds regression test ensuring wrapper descriptor AttributeErrors aren’t silently delegated.
tests/wrappers/test_autoreset_wrapper.py	Updates autoreset semantics: terminal fields remain top-level; final snapshot persists.
tests/wrappers/helpers.py	Adds helper envs and dtype/bounds tweaks used by new contract tests.
tests/test_struct.py	Updates static-field semantics and adds JAX-unflatten sentinel bypass tests.
tests/test_public_typing.py	New tests for exported typing names and py.typed marker.
tests/test_environment.py	New tests for abstract space descriptors and default_max_steps default.
tests/test_distribution_artifacts.py	New packaging tests for built wheel/sdist content boundaries and size limits.
tests/test_container.py	Tests canonical Container extras ordering and JAX control-flow composability.
tests/spaces/test_pytree_space.py	Loosens brittle error-message matching for structure mismatches.
tests/spaces/test_discrete.py	Strengthens shape/dtype guards; expands JIT/vmap coverage.
tests/spaces/test_continuous.py	Adjusts tests to reflect deferred invariant checks and stronger contains/sampling contracts.
tests/spaces/test_batched_space.py	Adds peel/rebatch round-trip tests and deferred batch-size validation coverage.
tests/adapters/test_navix_envelope.py	Removes xfails; validates entity cardinality and captured horizon warnings.
tests/adapters/test_mujoco_playground_envelope.py	Updates backend metadata placement and horizon pass-through warning behavior.
tests/adapters/test_kinetix_envelope.py	Adds level-id normalization tests and more robust horizon/kwargs behaviors.
tests/adapters/test_jumanji_envelope.py	Adds dtype preservation tests and explicit time-limit warning behavior.
tests/adapters/test_gymnax_regressions.py	New fake-backend Gymnax regression tests (horizon handling, backend info schema).
tests/adapters/test_create_integration.py	Refactors create-smoke tests to always init+step; removes max-steps override rejection test.
tests/adapters/test_craftax_envelope.py	Updates Craftax adapter behavior: horizon capture, backend kwargs overlap warnings, placeholders.
tests/adapters/test_common.py	New tests for shared adapter helper functions.
tests/adapters/test_brax_regressions.py	New fake-backend Brax regression tests (backend limits, done casting).
tests/adapters/test_brax_envelope.py	Updates Brax backend metadata placement and preserves pre-wrapped Brax envs.
src/envelope/wrappers/wrapper.py	Adds stack-constraint validation, marker mixins, safer attribute delegation, default_max_steps delegation.
src/envelope/wrappers/vmap_wrapper.py	Marks vmap wrapper as VectorizingWrapper.
src/envelope/wrappers/vmap_envs_wrapper.py	Marks env-vmap wrapper as VectorizingWrapper; makes spaces cached/consistent.
src/envelope/wrappers/truncation_wrapper.py	Enforces positive max_steps, preserves inner truncation, and updates init/reset semantics.
src/envelope/wrappers/state_injection_wrapper.py	Switches to full reset_info injection, adds active flag, and pooled-init exclusion constraint.
src/envelope/wrappers/pooled_init_vmap_wrapper.py	Adds type-preserving placeholders, final_valid, static config fields, and vectorization constraints.
src/envelope/wrappers/observation_normalization_wrapper.py	Adds shared-stats normalization for final, broadcast-aware stats reshaping, and space rebatching.
src/envelope/wrappers/normalization.py	Makes RMV count per-leaf and clamps variance for numerical safety.
src/envelope/wrappers/flatten_observation_wrapper.py	Preserves batch prefixes when flattening; uses rebatch helper.
src/envelope/wrappers/flatten_action_wrapper.py	Makes unflatten split indices static; supports batched actions cleanly.
src/envelope/wrappers/episode_statistics_wrapper.py	Resets episode stats on manual reset; zero-stats helper.
src/envelope/wrappers/continuous_observation_wrapper.py	Rebatches observation space and normalizes obs casting behavior.
src/envelope/wrappers/autoreset_wrapper.py	Introduces final_valid, type-preserving final placeholders, stack constraints, and implements reset.
src/envelope/wrappers/init.py	Re-exports stack constraint APIs and marker mixins.
src/envelope/typing.py	Exposes typing names via __all__; makes Info protocol read-only and update() return Self.
src/envelope/struct.py	Adds hashability validation for static fields, canonical Container extras ordering, and zeros_like.
src/envelope/spaces.py	Refines sampling/contains semantics; adds peel/rebatch helpers; makes contains return scalar JAX bool.
src/envelope/environment.py	Adds StackConstraint, helpers, and default_max_steps on base Environment.
src/envelope/adapters/navix_envelope.py	Captures/retains horizon, warns on overlap, fixes entity id cardinality, uses backend metadata container.
src/envelope/adapters/mujoco_playground_envelope.py	Captures default episode length, disables backend time limit, stores backend metadata under info.backend.
src/envelope/adapters/jumanji_envelope.py	Captures/disables default time_limit, normalizes boolean observations, moves extras into info.backend.
src/envelope/adapters/gymnax_envelope.py	Probes backend info schema, captures horizon, bypasses public auto-reset, stores backend info under info.backend.
src/envelope/adapters/craftax_envelope.py	Probes backend info schema, captures horizon, warns on overlap, stores backend info under info.backend.
src/envelope/adapters/brax_envelope.py	Captures/disables backend limits by default, warns on overlap, stores backend info under info.backend.
src/envelope/adapters/_common.py	New shared helpers for overlap warnings, horizon capture, backend container normalization, and schema probing.
src/envelope/adapters/init.py	Extends create() with max_episode_steps policy and uses default_max_steps.
src/envelope/init.py	Re-exports typing aliases at package root.
src/envelope/py.typed	Adds PEP 561 marker for typed wheels.
RELEASE_NOTES_0.5.md	New 0.5 release notes describing lifecycle, wrapper composition, spaces, adapters, and typing.
README.md	Updates wrapper composition messaging and installation guidance for adapter extras.
pyproject.toml	Bumps version to 0.5.0; adds optional deps/extras; tightens sdist include/exclude; adds ty/pre-commit config.
mkdocs.yml	Adjusts mkdocstrings options for stricter builds.
docs/api/wrappers.md	Documents stack constraints, ordering guidance, and new core wrapper infrastructure exports.
docs/api/struct.md	Documents canonical Container key ordering, static field hashability, and zeros_like.
docs/api/spaces.md	Documents peel/rebatch, deferred invariant checks, and stricter contains semantics.
docs/api/environments.md	Documents stack constraints and default_max_steps behavior.
docs/api/adapters.md	Documents adapter defaults, warnings, max_episode_steps, and info.backend metadata.
.pre-commit-config.yaml	Adds ruff lint/format hooks.
.gitignore	Adds RL output and local tool-selection ignores.
.github/workflows/publish.yml	Updates release workflow: matrix tests, build verification, artifact inspection, and wheel smoke test.
.github/workflows/ci.yml	New CI workflow: lint/format, tests, typing, min-jax, docs, and distribution checks.
.github/workflows/adapters.yml	New isolated adapter contract workflows for published extras and source-backed adapters.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.