Skip to content

Environment¤

jaxnasium.Environment ¤

Base environment class for JAX-compatible environments. Create your environment by subclassing this.

step and reset should typically not be overridden, as they merely handle the auto-reset logic. Instead, the environment-specific logic should be implemented in the step_env and reset_env methods.

step(key: PRNGKeyArray, state: TEnvState, action: PyTree[Real[Array, ...]]) -> tuple[TimeStep, TEnvState] ¤

Steps the environment forward with the given action and performs auto-reset when necessary. Additionally, this function inserts the original observation (before auto-resetting) in the info dictionary to bootstrap correctly on truncated episodes (info={"_TERMINAL_OBSERVATION": obs, ...})

This function should typically not be overridden. Instead, the environment-specific logic should be implemented in the step_env method.

Returns a TimeStep object (observation, reward, terminated, truncated, info) and the new state.

Arguments:

  • key: JAX PRNG key.
  • state: Current state of the environment.
  • action: Action to take in the environment.
step_env(key: PRNGKeyArray, state: TEnvState, action: PyTree[Real[Array, ...]]) -> tuple[TimeStep, TEnvState] abstractmethod ¤

Defines the environment-specific step logic. I.e. here the state of the environment is updated according to the transition function.

Returns a TimeStep object (observation, reward, terminated, truncated, info) and the new state.

Arguments:

  • key: JAX PRNG key.
  • state: Current state of the environment.
  • action: Action to take in the environment.
reset(key: PRNGKeyArray) -> tuple[Observation, TEnvState] ¤

Resets the environment to an initial state and returns the initial observation. Environment-specific logic is defined in the reset_env method. Typically, this function should not be overridden.

Returns the initial observation and the initial state of the environment.

Arguments:

  • key: JAX PRNG key.
reset_env(key: PRNGKeyArray) -> tuple[Observation, TEnvState] abstractmethod ¤

Defines the environment-specific reset logic.

Returns the initial observation and the initial state of the environment.

Arguments:

  • key: JAX PRNG key.
observation_space: Space | PyTree[Space] abstractmethod property ¤

Defines the space of possible observations from the environment. For multi-agent environments, this should be a PyTree of spaces. See jaxnasium.spaces for more information on how to define (composite) observation spaces.

sample_observation(key: PRNGKeyArray) -> Observation ¤

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

action_space: Space | PyTree[Space] abstractmethod property ¤

Defines the space of valid actions for the environment. For multi-agent environments, this should be a PyTree of spaces. See jaxnasium.spaces for more information on how to define (composite) action spaces.

sample_action(key: PRNGKeyArray) -> PyTree[Real[Array, ...]] ¤

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

auto_reset(key: PRNGKeyArray, timestep_step: TimeStep, state_step: TEnvState) -> tuple[TimeStep, TEnvState] ¤

Auto-resets the environment when the episode is terminated or truncated.

Given a step timestep and state, this function will auto-reset the environment and return the new timestep and state when the episode is terminated or truncated. Inserts the original observation in info to bootstrap correctly on truncated episodes.

Arguments:

  • key: JAX PRNG key.
  • timestep_step: The timestep returned by the step_env method.
  • state_step: The state returned by the step_env method.

Returns: A tuple of the new timestep and state with the state and observation reset to a new initial state and observation when the episode is terminated or truncated. The original observation is inserted in info to bootstrap correctly on truncated episodes.

multi_agent: bool property ¤

Indicates if the environment is a multi-agent environment.

Infers this via the _multi_agent property. If not set, assumes single-agent.

agent_structure: PyTreeDef property ¤

Returns the structure of the agent space. In single-agent environments this is simply a PyTreeDef(). However, for multi-agent environments this is a PyTreeDef((, x num_agents)).

state_type: type[TEnvState] property ¤

Optional property for accessing the class of the environment state, if set through the type parameter of this environment. That is, the environment (or one of its parent classes) declares its state type via class Env(Environment[MyState]), or through the default of a TypeVar. If this is not the case, this raises a TypeError.

Generally not required, but possibly usefull when an environment needs and extensible state.

Note

When subclassing this class with a type parameter (Environment[MyState]), the state type is available at runtime via state_type; this may be helpfull for environments that should have an extensible state. The state type is read from class definitions only: instantiating a subscripted class, e.g. MyEnv[ExtendedState](), silently creates a MyEnv with its default state type. Declare a subclass instead.

Timestep¤

jaxnasium.TimeStep ¤

A container for the output of an environment's step function. (timestep, state = env.step(...)).

This class follows the Gymnasium standard API, with the signature: (obs, reward, terminated, truncated, info) tuple.

Arguments:

  • observation: The environment state representation provided to the agent. Can be an Array or a PyTree of arrays. When using action masking, the observation should be of type AgentObservation.
  • reward: The reward signal from the previous action, indicating performance. Can be a scalar Array or a PyTree of reward Arrays (in the case of multi agent-environments).
  • terminated: Boolean flag indicating whether the episode has ended due to reaching a terminal state (e.g., goal reached, game over).
  • truncated: Boolean flag indicating whether the episode ended due to external factors (e.g., reaching max steps, timeout).
  • info: Dictionary containing any additional information about the environment step.

Environment state¤

jaxnasium.EnvState ¤

A simple container for the state of an environment. Defined as an Equinox Module (a PyTree registered dataclass) with a helper method for replacing fields.

It is not required to use this class for the state of an environment.

Observation container (optional)¤

jaxnasium.AgentObservation ¤

A container for the observation of a single agent, with optional action masking.

Typically, this container is optional. However, Algorithms in jaxnasium.algorithms expect observations to be wrapped in this type when action masking is enabled.

Arguments:

  • observation: The observation of the agent.
  • action_mask: The action mask of the agent. A boolean array of the same shape as the action space.
  • critic_observation: Optional observation of the critic for the agent. when not provided, the regular observation is used. This may be used for CTDE (like) approaches.