Skip to content

Adapters

Creating adapters

Envelope has a powerful create method that can create environments from many different JAX-based environment suites, and automatically casts them into the envelope Environment interface. The create function uses a unique environment identifier of the format SUITE::ENVIRONMENT, and can pass arguments to the environment constructor:

env = create("brax::ant", env_kwargs={"backend": "positional"})

Note: Many suites natively support episode truncation and automatic resetting. When instantiating adapters, users are highly discouraged from using these features. Instead, they should use the envelope-native wrapper ecosystem.

Adapters implement the HasFromNameInit protocol, that ensures they can be created via a from_name(cls, env_name: str, env_kwargs=None, **kwargs) function. This is used by create. The env_kwargs are passed to the suite's environment constructor (such as navix.make or brax.envs.create), and the kwargs are passed to the adapter class on creation.

Each adapter may have a default_episode_length property that is populated depending on the suite and specific environment. If it is not None, the create function will wrap the adapter in a TruncationWrapper before returning it.

API Reference (create)

envelope.adapters.create(env_name, env_kwargs=None, **kwargs)

Create an environment from a prefixed environment ID.

Parameters:

Name Type Description Default
env_name str

Environment ID in the format "suite::env_name" (e.g., "brax::ant")

required
env_kwargs dict[str, Any] | None

Keyword arguments passed to the suite's environment constructor

None
**kwargs dict[str, Any]

Additional keyword arguments passed to the environment wrapper

{}

Returns:

Type Description
Environment

An instance of the wrapped environment. If the adapter has a default_max_steps property, it will be wrapped in a TruncationWrapper before returning.

Examples:

>>> env = create("jumanji::snake")
>>> env = create("brax::ant", env_kwargs={"backend": "spring"})
>>> env = create("gymnax::CartPole-v1", env_params=...)

envelope.adapters.HasFromNameInit

Bases: Protocol

Functions

from_name(env_name, env_kwargs=None, **kwargs) classmethod

Creates an environment from a name and keyword arguments. Unless otherwise noted, the created environment will have its default parameters, with truncation and auto-reset disabled.

Parameters:

Name Type Description Default
env_name str

Environment name

required
env_kwargs dict[str, Any] | None

Keyword arguments passed to the environment constructor

None
**kwargs dict[str, Any]

Additional keyword arguments passed to the environment wrapper

{}

API Reference (Specific Adapters)

envelope.adapters.brax_envelope.BraxEnvelope

Bases: Environment

Wrapper to convert a Brax environment to a envelope environment.

Note that Brax' create function has a default value of 1000 for the episode length. Thus, the default_max_steps property is set to 1000 for all Brax envs.

Brax uses a dataclass as the environment state, which we return in full as fields of the Info object.

Parameters:

Name Type Description Default
brax_env Env

the Brax environment.

required

Functions

from_name(env_name, env_kwargs=None) classmethod

Create a BraxEnvelope from a name and keyword arguments. env_kwargs arepassed to brax.envs.create.

envelope.adapters.craftax_envelope.CraftaxEnvelope

Bases: Environment

Wrapper to convert a Craftax environment to a envelope environment.

Craftax (mostly) uses the Gymnax interface, so the info object is only created on the first step. To keep structural equivalence between the init and step infos, we create a placeholder filled with jnp.nan that is returned on init.

Parameters:

Name Type Description Default
craftax_env CraftaxEnv

the Craftax environment, with baked-in static_env_params.

required
env_params CraftaxEnvParams

the environment parameters, which are passed to the Craftax environment's reset and step methods.

required

Functions

from_name(env_name, env_params=None, env_kwargs=None) classmethod

Create a CraftaxEnvelope from a name and keyword arguments. env_kwargs are passed to craftax.craftax_env.make_craftax_env_from_name.

envelope.adapters.gymnax_envelope.GymnaxEnvelope

Bases: Environment

Wrapper to convert a Gymnax environment to a envelope environment.

Gymnax interface only creates the info object on the first step. To keep structural equivalence between the init and step infos, we create a placeholder filled with jnp.nan that is returned on init.

Gymnax implements Tuple and Dict spaces, which are converted to PyTreeSpace of a tuple and dict PyTree respectively.

Parameters:

Name Type Description Default
gymnax_env Environment

the Gymnax environment.

required
env_params EnvParams

the environment parameters, which are passed to the Gymnax environment's reset and step methods.

required

Functions

from_name(env_name, env_params=None, env_kwargs=None) classmethod

Create a GymnaxEnvelope from a name and keyword arguments. env_kwargs are passed to gymnax.make.

envelope.adapters.jumanji_envelope.JumanjiEnvelope

Bases: Environment

Wrapper to convert a Jumanji environment to a envelope environment.

Some Jumanji environments support time limits via a time_limit attribute of the environemnt. If this attribute exists, we overwrite it with the maximum integer value and set default_max_steps to the original value.

Parameters:

Name Type Description Default
jumanji_env Environment

the Jumanji environment.

required

Functions

from_name(env_name, env_kwargs=None) classmethod

Create a JumanjiEnvelope from a name and keyword arguments. env_kwargs are passed to jumanji.make.

envelope.adapters.kinetix_envelope.KinetixEnvelope

Bases: Environment

Wrapper to convert a Kinetix environment to an envelope environment.

Kinetix environments are constructed via a reset_fn that produces a level on each reset, rather than a simple environment name. Two creation modes are provided: create_from_size and create_random.

Kinetix only produces the env_info dict on the first step, not on reset. To keep structural equivalence between the init and step infos (required for jax.lax.scan, jax.vmap, etc.), a NaN-filled placeholder with the same pytree structure is returned on init.

Parameters:

Name Type Description Default
kinetix_env KinetixEnv

the Kinetix environment, with baked-in static_env_params.

required
env_params EnvParams

the environment parameters, which are passed to the Kinetix environment's reset and step methods.

required

Functions

create_from_size(size, action_type=ActionType.CONTINUOUS, observation_type=ObservationType.SYMBOLIC_FLAT, auto_reset=False) classmethod

Reset to a random level from a size category on each reset.

Loads all packaged levels for the given size ("s", "m", "l") and builds a reset function that uniformly samples one on each call.

create_random(action_type=ActionType.CONTINUOUS, observation_type=ObservationType.SYMBOLIC_FLAT, env_params=None, static_env_params=KinetixStaticEnvParams(), auto_reset=False) classmethod

Create a random level on each reset using Kinetix's kinetix.environment.ued.ued.make_reset_fn_sample_kinetix_level.

from_name(env_name, env_kwargs=None) classmethod

Dispatch to the appropriate creation mode.

  • "s", "m", "l": reset to a random level from that size category via create_from_size.
  • "random": fully procedural UED levels via create_random.

envelope.adapters.mujoco_playground_envelope.MujocoPlaygroundEnvelope

Bases: Environment

Wrapper to convert a mujoco_playground environment to a envelope environment.

Mujoco Playground uses a dataclass as the environment state, which we return in full as fields of the Info object.

All Mujoco Playground environments have continuous actions and observations, which range between (-1, 1) and (-inf, inf) respectively.

Parameters:

Name Type Description Default
mujoco_playground_env MjxEnv

the Mujoco Playground environment.

required

Functions

from_name(env_name, env_kwargs=None) classmethod

Create a MujocoPlaygroundEnvelope from a name and keyword arguments. env_kwargs are passed to config_overrides of mujoco_playground.registry.load.

envelope.adapters.navix_envelope.NavixEnvelope

Bases: Environment

Wrapper to convert a Navix environment to a envelope environment.

Navix uses a dataclass as the environment state, which we return in full as fields of the Info object.

Parameters:

Name Type Description Default
navix_env Environment

the Navix environment.

required

Functions

from_name(env_name, env_kwargs=None) classmethod

Create a NavixEnvelope from a name and keyword arguments. env_kwargs are passed to navix.make.