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 |
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
|
required |
env_params
|
CraftaxEnvParams
|
the environment parameters, which are passed to
the Craftax environment's |
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 |
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
|
required |
env_params
|
EnvParams
|
the environment parameters, which are passed
to the Kinetix environment's |
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 viacreate_from_size."random": fully procedural UED levels viacreate_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.