Skip to content

Spaces¤

Spaces describe the shape and range of observations and actions. Jaxnsium environments may return a single Space, or a PyTree of spaces (Additionally, see Multi-Agent).

jaxnasium.Space ¤

The base class for all spaces in Jaxnasium. Instead of using this class directly, use the subclasses Box, Discrete, and MultiDiscrete. Composite spaces can be created by simply combining spaces in an arbitrary PyTree. For example, a tuple of Box spaces can be created as follows: 

from jaxnasium import Box

box1 = Box(low=0, high=1, shape=(3,))
box2 = Box(low=0, high=1, shape=(4,))
box3 = Box(low=0, high=1, shape=(5,))
composite_space = (box1, box2, box3)

Jaxnasium algorithms assume multi-agent environments are such a composite space, where the first level of the PyTree is the agent dimension. For example, a multi-agent environment observation space may look like this: 

from jaxnasium import Box
from jaxnasium import MultiDiscrete

agent1_obs = Box(low=0, high=1, shape=(3,))
agent2_obs = Box(low=0, high=1, shape=(4,))
agent3_obs = MultiDiscrete(nvec=[2, 3])
env_obs_space = {
    "agent1": agent1_obs,
    "agent2": agent2_obs,
    "agent3": agent3_obs,
}

Spaces are purposefully not registered PyTree nodes.

shape: eqx.AbstractVar[tuple[int, ...]] instance-attribute ¤
dtype: eqx.AbstractVar[DTypeLike] instance-attribute ¤
sample(rng: PRNGKeyArray) -> Array abstractmethod ¤

jaxnasium.Box dataclass ¤

The standard Box space for continuous action/observation spaces.

Arguments:

  • low (int / Array[int]): The lower bound of the space.
  • high (int / Array[int]): The upper bound of the space.
  • shape: The shape of the space.
  • dtype: The data type of the space. Default is jnp.float32.
sample(rng: PRNGKeyArray) -> Array ¤

Sample a random element of the space.

As is the case in Gymnasium, floating dtypes are sampled differently based on the bounds of this space. If fully bounded, we sample uniformly, unbounded spaces are sampled with a gaussian, and spaces bounded from below or above are sampled from 'low + exponential' or 'high - exponential' respectively

jaxnasium.Discrete dataclass ¤

The standard discrete space for discrete action/observation spaces.

Arguments:

  • n (int): The number of discrete actions.
  • dtype: The data type of the space. Default is jnp.int16.
sample(rng: PRNGKeyArray) -> Int[Array, ''] ¤

Sample random action uniformly from set of discrete choices.

jaxnasium.MultiDiscrete dataclass ¤

The standard multi-discrete space for discrete action/observation spaces.

Arguments:

  • nvec (Array[int]): The number of discrete actions for each dimension.
  • dtype: The data type of the space. Default is jnp.int32.
sample(rng: PRNGKeyArray) -> Int[Array, ''] ¤

Sample random action uniformly from set of discrete choices.