Skip to content

unienv_interface

Env

Bases: ABC, Generic[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Stateful environment interface used throughout UniEnv.

An Env exposes typed action, observation, and optional context spaces plus the standard reset/step/render lifecycle. Concrete implementations may be single-instance or batched depending on batch_size.

metadata class-attribute instance-attribute

metadata: dict[str, Any] = {'render_modes': []}

render_mode class-attribute instance-attribute

render_mode: Optional[str] = None

render_fps class-attribute instance-attribute

render_fps: Optional[int] = None

backend instance-attribute

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device instance-attribute

device: Optional[BDeviceType]

batch_size class-attribute instance-attribute

batch_size: Optional[int] = None

action_space instance-attribute

action_space: Space[ActType, BDeviceType, BDtypeType, BRNGType]

observation_space instance-attribute

observation_space: Space[ObsType, BDeviceType, BDtypeType, BRNGType]

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

rng class-attribute instance-attribute

rng: BRNGType = None

unwrapped property

unwrapped: Env

prev_wrapper_layer property

prev_wrapper_layer: Optional[Env]

step abstractmethod

step(action: ActType) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

Advance the environment by one control step.

reset abstractmethod

reset(*, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]

Resets the environment to its initial state and returns the initial context and observation. If mask is provided, it will only return the masked context and observation, so the batch dimension in the output will not be same as the batch dimension in the context and observation spaces.

render abstractmethod

render() -> RenderFrame | Sequence[RenderFrame] | None

Render the current environment state using the configured render mode.

close

close()

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

SyncVecEnv

SyncVecEnv(env_fn: Iterable[Callable[[], Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]]], seed: Optional[int] = None)

Bases: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Run multiple non-batched environments synchronously in one process.

Instantiate every environment immediately and batch their interfaces.

unwrapped property

unwrapped: Env

prev_wrapper_layer property

prev_wrapper_layer: Optional[Env]

envs instance-attribute

envs = [(fn()) for fn in env_fn]

action_space instance-attribute

action_space = batch_differing_spaces([(action_space) for env in (envs)], device=device)

observation_space instance-attribute

observation_space = batch_differing_spaces([(observation_space) for env in (envs)], device=device)

context_space instance-attribute

context_space = None if context_space is None else batch_differing_spaces([(context_space) for env in (envs)], device=device)

rng instance-attribute

rng = random_number_generator(seed, device=device)

metadata property

metadata: Dict[str, Any]

render_mode property

render_mode: Optional[str]

render_fps property

render_fps: Optional[int]

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device property

device: Optional[BDeviceType]

batch_size property

batch_size: Optional[int]

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

reset

reset(*args, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]

Reset all selected workers and concatenate their results.

step

step(action: ActType) -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]

Step every worker environment once and concatenate the outputs.

render

render() -> Sequence[RenderFrame] | None

Collect render outputs from every worker that returns a frame.

close

close()

AsyncVecEnv

AsyncVecEnv(env_fn: Iterable[Callable[[], Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]]], seed: Optional[int] = None, ctx: Optional[BaseContext] = None, daemon: bool = True)

Bases: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Run multiple non-batched environments in separate worker processes.

Spawn worker processes and batch their exposed spaces.

unwrapped property

unwrapped: Env

prev_wrapper_layer property

prev_wrapper_layer: Optional[Env]

command_pipes instance-attribute

command_pipes: List[Connection] = []

processes instance-attribute

processes: List[Process] = []

error_queue instance-attribute

error_queue: Queue = Queue()

backend instance-attribute

backend = backend

device instance-attribute

device = device

metadata instance-attribute

metadata = metadata

render_mode instance-attribute

render_mode = render_mode

render_fps instance-attribute

render_fps = render_fps

action_space instance-attribute

action_space = batch_space(action_space, len(processes))

observation_space instance-attribute

observation_space = batch_space(observation_space, len(processes))

context_space instance-attribute

context_space = None if context_space is None else batch_space(context_space, len(processes))

rng instance-attribute

rng = random_number_generator(seed, device=device)

batch_size property

batch_size: Optional[int]

render abstractmethod

render() -> RenderFrame | Sequence[RenderFrame] | None

Render the current environment state using the configured render mode.

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

send_command

send_command(index: int, cmd: Literal['reset', 'step', 'render', 'close'], *args, **kwargs)

Send a command to one worker process.

get_command_result

get_command_result(index: int)

Receive one worker result and surface worker-side exceptions.

reset

reset(*args, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]

Convenience wrapper around reset_async plus reset_wait.

reset_async

reset_async(*args, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> None

Dispatch reset commands to all selected workers.

reset_wait

reset_wait() -> Tuple[ContextType, ObsType, Dict[str, Any]]

Collect the pending reset results and batch them.

step

step(action: ActType) -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]

Convenience wrapper around step_async plus step_wait.

step_async

step_async(action: ActType) -> None

Dispatch one action batch to all worker processes.

step_wait

step_wait() -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]

Collect the pending step results and batch them.

close

close()

Shut down all worker processes and communication pipes.

Wrapper

Wrapper(env: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType])

Bases: Env[WrapperBArrayT, WrapperContextT, WrapperObsT, WrapperActT, WrapperRenderFrame, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT], Generic[WrapperBArrayT, WrapperContextT, WrapperObsT, WrapperActT, WrapperRenderFrame, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Base class for stateful environment wrappers.

The wrapper forwards all operations to env by default while allowing subclasses to override spaces, RNG handling, and selected lifecycle methods.

Wrap env and optionally override spaces or metadata later.

metadata property writable

metadata: Dict[str, Any]

Returns the :attr:Env :attr:metadata.

render_mode property

render_mode: Optional[WrapperRenderFrame]

render_fps property

render_fps: Optional[int]

backend property

backend: ComputeBackend[Any, WrapperBDeviceT, Any, WrapperBRngT]

device property

device: Optional[WrapperBDeviceT]

batch_size property

batch_size: Optional[int]

action_space property writable

action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]

observation_space property writable

observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]

context_space property writable

context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]

rng property writable

rng: WrapperBRngT

Returns the :attr:Env :attr:rng attribute.

env instance-attribute

env = env

unwrapped property

unwrapped: Env

Returns the base environment of the wrapper.

This will be the bare :class:gymnasium.Env environment, underneath all layers of wrappers.

prev_wrapper_layer property

prev_wrapper_layer: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

step

step(action: WrapperActT) -> Tuple[WrapperObsT, Union[SupportsFloat, WrapperBArrayT], Union[bool, WrapperBArrayT], Union[bool, WrapperBArrayT], Dict[str, Any]]

reset

reset(*args, mask: Optional[WrapperBArrayT] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]

render

render() -> RenderFrame | Sequence[RenderFrame] | None

close

close()

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the given attribute is within the wrapper or its environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets an attribute from the wrapper and lower environments if name doesn't exist in this object.

Parameters:

Name Type Description Default
name str

The variable name to get

required

Returns:

Type Description
Any

The variable with name in wrapper or lower environments

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets an attribute on this wrapper or lower environment if name is already defined.

Parameters:

Name Type Description Default
name str

The variable name

required
value Any

The new variable value

required

ActionWrapper

ActionWrapper(env: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType])

Bases: Wrapper[BArrayType, ContextType, ObsType, WrapperActT, RenderFrame, BDeviceType, BDtypeType, BRNGType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType], Generic[WrapperActT, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Wrapper that exposes a transformed action interface.

Wrap env and optionally override spaces or metadata later.

metadata property writable

metadata: Dict[str, Any]

Returns the :attr:Env :attr:metadata.

render_mode property

render_mode: Optional[WrapperRenderFrame]

render_fps property

render_fps: Optional[int]

backend property

backend: ComputeBackend[Any, WrapperBDeviceT, Any, WrapperBRngT]

device property

device: Optional[WrapperBDeviceT]

batch_size property

batch_size: Optional[int]

action_space property writable

action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]

observation_space property writable

observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]

context_space property writable

context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]

rng property writable

rng: WrapperBRngT

Returns the :attr:Env :attr:rng attribute.

unwrapped property

unwrapped: Env

Returns the base environment of the wrapper.

This will be the bare :class:gymnasium.Env environment, underneath all layers of wrappers.

prev_wrapper_layer property

prev_wrapper_layer: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

env instance-attribute

env = env

reset

reset(*args, mask: Optional[WrapperBArrayT] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]

render

render() -> RenderFrame | Sequence[RenderFrame] | None

close

close()

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the given attribute is within the wrapper or its environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets an attribute from the wrapper and lower environments if name doesn't exist in this object.

Parameters:

Name Type Description Default
name str

The variable name to get

required

Returns:

Type Description
Any

The variable with name in wrapper or lower environments

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets an attribute on this wrapper or lower environment if name is already defined.

Parameters:

Name Type Description Default
name str

The variable name

required
value Any

The new variable value

required

map_action abstractmethod

map_action(action: WrapperActT) -> ActType

Map an outer action into the wrapped environment's action space.

reverse_map_action

reverse_map_action(action: ActType) -> WrapperActT

Map an inner action back to the wrapper-facing representation.

step

step(action: WrapperActT) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

ContextObservationWrapper

ContextObservationWrapper(env: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType])

Bases: Wrapper[BArrayType, WrapperContextT, WrapperObsT, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType], Generic[WrapperContextT, WrapperObsT, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Wrapper that transforms context and observation values.

Wrap env and optionally override spaces or metadata later.

metadata property writable

metadata: Dict[str, Any]

Returns the :attr:Env :attr:metadata.

render_mode property

render_mode: Optional[WrapperRenderFrame]

render_fps property

render_fps: Optional[int]

backend property

backend: ComputeBackend[Any, WrapperBDeviceT, Any, WrapperBRngT]

device property

device: Optional[WrapperBDeviceT]

batch_size property

batch_size: Optional[int]

action_space property writable

action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]

observation_space property writable

observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]

context_space property writable

context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]

rng property writable

rng: WrapperBRngT

Returns the :attr:Env :attr:rng attribute.

unwrapped property

unwrapped: Env

Returns the base environment of the wrapper.

This will be the bare :class:gymnasium.Env environment, underneath all layers of wrappers.

prev_wrapper_layer property

prev_wrapper_layer: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

env instance-attribute

env = env

render

render() -> RenderFrame | Sequence[RenderFrame] | None

close

close()

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the given attribute is within the wrapper or its environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets an attribute from the wrapper and lower environments if name doesn't exist in this object.

Parameters:

Name Type Description Default
name str

The variable name to get

required

Returns:

Type Description
Any

The variable with name in wrapper or lower environments

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets an attribute on this wrapper or lower environment if name is already defined.

Parameters:

Name Type Description Default
name str

The variable name

required
value Any

The new variable value

required

map_context

map_context(context: WrapperContextT) -> ContextType

Convert wrapped context into the wrapper-facing representation.

reverse_map_context

reverse_map_context(context: ContextType) -> WrapperContextT

Convert wrapper-facing context back into the wrapped representation.

map_observation

map_observation(observation: WrapperObsT) -> ObsType

Convert wrapped observations into the wrapper-facing representation.

reverse_map_observation

reverse_map_observation(observation: ObsType) -> WrapperObsT

Convert wrapper-facing observations back into the wrapped representation.

reset

reset(*args, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]

step

step(action: ActType) -> Tuple[WrapperObsT, Union[float, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

FuncEnv

Bases: ABC, Generic[StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Functional environment interface with explicit immutable state passing.

Unlike :class:Env, a FuncEnv never stores simulation state on the instance. Instead, every lifecycle method takes a state object and returns the updated state explicitly, which makes the interface easier to compose with JAX-style functional code.

metadata class-attribute instance-attribute

metadata: Dict[str, Any] = {'render_modes': []}

backend instance-attribute

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device class-attribute instance-attribute

device: Optional[BDeviceType] = None

batch_size class-attribute instance-attribute

batch_size: Optional[int] = None

observation_space instance-attribute

observation_space: Space[Any, BDeviceType, BDtypeType, BRNGType]

action_space instance-attribute

action_space: Space[Any, BDeviceType, BDtypeType, BRNGType]

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

unwrapped property

unwrapped: FuncEnv

prev_wrapper_layer property

prev_wrapper_layer: Optional[FuncEnv]

initial abstractmethod

initial(*, seed: Optional[int] = None, **kwargs) -> Tuple[StateType, ContextType, ObsType, Dict[str, Any]]

Create the first environment state, context, observation, and info.

reset abstractmethod

reset(state: StateType, *, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> Tuple[StateType, ContextType, ObsType, Dict[str, Any]]

Resets the environment to its initial state and returns the initial context and observation. If mask is provided, it will only return the masked context and observation, so the batch dimension in the output will not be same as the batch dimension in the context and observation spaces. Note that state input and output should be with full batch dimensions

step abstractmethod

step(state: StateType, action: ActType) -> Tuple[StateType, ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

Advance state by one control step using action.

close

close(state: StateType) -> None

Release any resources associated with state.

render_init

render_init(state: StateType, *, seed: Optional[int] = None, render_mode: Optional[str] = None, **kwargs) -> Tuple[StateType, RenderStateType, FuncEnvCommonRenderInfo]

Create any auxiliary render state needed by render_image.

render_image

render_image(state: StateType, render_state: RenderStateType) -> Tuple[RenderFrame | Sequence[RenderFrame] | None, StateType, RenderStateType]

Produce a render frame and updated runtime state.

render_close

render_close(state: StateType, render_state: RenderStateType) -> StateType

Release render-specific resources and return the updated state.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

FuncEnvBasedEnv

FuncEnvBasedEnv(func_env: FuncEnv[StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType], *, render_mode: Optional[str] = None, render_kwargs: Dict[str, Any] = {})

Bases: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType], Generic[StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Adapter that exposes a stateful Env view over a FuncEnv.

Wrap a functional environment and manage its runtime state internally.

unwrapped property

unwrapped: Env

prev_wrapper_layer property

prev_wrapper_layer: Optional[Env]

func_env instance-attribute

func_env = func_env

state instance-attribute

state: Optional[StateType] = None

render_state instance-attribute

render_state: Optional[RenderStateType] = None

rng instance-attribute

rng = random_number_generator(device=device)

metadata property writable

metadata: Dict[str, Any]

render_mode property

render_mode: Optional[str]

render_fps property

render_fps: Optional[int]

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device property

device: Optional[BDeviceType]

batch_size property

batch_size: Optional[int]

action_space property

action_space: Space[ActType, BDeviceType, BDtypeType, BRNGType]

observation_space property

observation_space: Space[ObsType, BDeviceType, BDtypeType, BRNGType]

context_space property

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]]

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

step

step(action: ActType) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

reset

reset(*, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]

render

render() -> RenderFrame | Sequence[RenderFrame] | None

close

close() -> None

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

FuncEnvWrapper

FuncEnvWrapper(func_env: FuncEnv[StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType])

Bases: Generic[WrapperStateT, WrapperRenderStateT, WrapperBArrayT, WrapperContextT, WrapperObsT, WrapperActT, WrapperRenderFrame, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT, StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType], FuncEnv[WrapperStateT, WrapperRenderStateT, WrapperBArrayT, WrapperContextT, WrapperObsT, WrapperActT, WrapperRenderFrame, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]

Base wrapper for FuncEnv implementations.

This mirrors :class:Wrapper but preserves the explicit state-passing style of the functional environment API.

Wrap a functional environment and optionally override its interface metadata.

metadata property writable

metadata: Dict[str, Any]

Returns the :attr:Env :attr:metadata.

backend property

backend: ComputeBackend[Any, WrapperBDeviceT, Any, WrapperBRngT]

device property

device: Optional[WrapperBDeviceT]

batch_size property

batch_size: Optional[int]

action_space property writable

action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]

observation_space property writable

observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]

context_space property writable

context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]

func_env instance-attribute

func_env = func_env

unwrapped property

unwrapped: FuncEnv

prev_wrapper_layer property

prev_wrapper_layer: FuncEnv[StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

initial

initial(*, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperStateT, WrapperContextT, WrapperObsT, Dict[str, Any]]

reset

reset(state: WrapperStateT, *args, seed: Optional[int] = None, mask: Optional[WrapperBArrayT] = None, **kwargs) -> Tuple[WrapperStateT, WrapperContextT, WrapperObsT, Dict[str, Any]]

step

step(state: WrapperStateT, action: WrapperActT) -> Tuple[WrapperStateT, WrapperObsT, Union[SupportsFloat, WrapperBArrayT], Union[bool, WrapperBArrayT], Union[bool, WrapperBArrayT], Dict[str, Any]]

close

close(state: WrapperStateT) -> None

render_init

render_init(state: WrapperStateT, *, seed: Optional[int] = None, render_mode: Optional[str] = None, **kwargs) -> Tuple[WrapperStateT, WrapperRenderStateT, FuncEnvCommonRenderInfo]

render_image

render_image(state: WrapperStateT, render_state: WrapperRenderStateT) -> Tuple[WrapperRenderFrame | Sequence[WrapperRenderFrame] | None, WrapperStateT, WrapperRenderStateT]

render_close

render_close(state: WrapperStateT, render_state: WrapperRenderStateT) -> WrapperStateT

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

World

Bases: ABC, Generic[BArrayType, BDeviceType, BDtypeType, BRNGType]

Mutable world/simulator interface used by node-based environments.

A World owns the shared simulation state while WorldNode instances contribute observations, actions, rewards, and reset/reload logic around it.

backend instance-attribute

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device instance-attribute

device: Optional[BDeviceType]

The world timestep in seconds, if None, the world is asynchronous (real-time)

world_timestep instance-attribute

world_timestep: Optional[float]

The world's physical timestep in seconds, there might be multiple world sub-steps inside a world step. If none this means it is not known

world_subtimestep class-attribute instance-attribute

world_subtimestep: Optional[float] = None

The number of parallel environments in this world

batch_size class-attribute instance-attribute

batch_size: Optional[int] = None

step abstractmethod

step() -> Union[float, BArrayType]

Step the world by one timestep. Returns: float: The elapsed time since the last step.

reset abstractmethod

reset(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Reset the world state in-place.

reload

reload(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Prepare the world for reloading. This is called BEFORE nodes add entities.

Use this hook to: - Clear old scene data - Load asset definitions - Prepare for entities to be added by nodes

DO NOT compile the scene here. Compilation happens in after_reload() after all nodes have had a chance to add their entities.

By default, this just calls reset with the same parameters. Simulation environments should override this to prepare for entity addition.

The reload flow is
  1. world.reload() # World prepares (clear, load assets)
  2. node.reload() # Nodes add entities to scene
  3. world.after_reload() # World compiles scene with all entities
  4. node.after_reload() # Nodes cache references

after_reset

after_reset(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Called after reset() completes.

Use this hook to perform post-reset initialization that requires the simulation to be in a valid state (e.g., caching entity indices, validating scene state).

This is called by the environment composer after all nodes have been reset.

after_reload

after_reload(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Compile the world scene AFTER all nodes have added their entities.

This is where the actual scene compilation should happen, since nodes have already added their entities during the reload() phase.

Use this hook to: - Compile the simulation scene - Cache references to entities - Perform validation that requires the compiled scene - Initialize state that depends on the final scene configuration

This is called by the environment composer after all nodes have been reloaded.

close

close() -> None

Release any resources owned by the world.

is_control_timestep_compatible

is_control_timestep_compatible(control_timestep: Optional[float]) -> bool

Check whether a node control period aligns with world_timestep.

RealWorld

RealWorld(backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], device: Optional[BDeviceType] = None, world_timestep: Optional[float] = None, world_subtimestep: Optional[float] = None, batch_size: Optional[int] = None)

Bases: World[BArrayType, BDeviceType, BDtypeType, BRNGType]

World implementation backed by wall-clock elapsed time.

backend instance-attribute

backend = backend

device instance-attribute

device = device

world_timestep instance-attribute

world_timestep = world_timestep

world_subtimestep instance-attribute

world_subtimestep = world_subtimestep

batch_size instance-attribute

batch_size = batch_size

reload

reload(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Prepare the world for reloading. This is called BEFORE nodes add entities.

Use this hook to: - Clear old scene data - Load asset definitions - Prepare for entities to be added by nodes

DO NOT compile the scene here. Compilation happens in after_reload() after all nodes have had a chance to add their entities.

By default, this just calls reset with the same parameters. Simulation environments should override this to prepare for entity addition.

The reload flow is
  1. world.reload() # World prepares (clear, load assets)
  2. node.reload() # Nodes add entities to scene
  3. world.after_reload() # World compiles scene with all entities
  4. node.after_reload() # Nodes cache references

after_reset

after_reset(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Called after reset() completes.

Use this hook to perform post-reset initialization that requires the simulation to be in a valid state (e.g., caching entity indices, validating scene state).

This is called by the environment composer after all nodes have been reset.

after_reload

after_reload(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Compile the world scene AFTER all nodes have added their entities.

This is where the actual scene compilation should happen, since nodes have already added their entities during the reload() phase.

Use this hook to: - Compile the simulation scene - Cache references to entities - Perform validation that requires the compiled scene - Initialize state that depends on the final scene configuration

This is called by the environment composer after all nodes have been reloaded.

close

close() -> None

Release any resources owned by the world.

is_control_timestep_compatible

is_control_timestep_compatible(control_timestep: Optional[float]) -> bool

Check whether a node control period aligns with world_timestep.

step

step() -> float

Return the wall-clock delta since the previous step.

reset

reset() -> None

Start a new wall-clock timing sequence.

WorldNode

Bases: ABC, Generic[ContextType, ObsType, ActType, BArrayType, BDeviceType, BDtypeType, BRNGType]

A stateful node that manages one aspect of a world (e.g. a sensor, a robot, a reward function).

The environment is initialized during the first call to reset or reload.

Lifecycle — reset flow::

World.reset()
  -> WorldNode.reset(priority=...)
  -> WorldNode.after_reset(priority=...)
  -> WorldNode.get_context() / get_observation() / get_info() / render()

Lifecycle — reload flow::

World.reload()
  -> WorldNode.reload(priority=...)
  -> WorldNode.after_reset(priority=...)
  -> WorldNode.get_context() / get_observation() / get_info() / render()

reload re-generates the simulation environment (e.g. re-reading assets, rebuilding the scene). This is typically much more expensive than reset and should only be called when the environment configuration has changed. By default reload delegates to reset.

Lifecycle — step flow::

WorldNode.set_next_action(action)
  -> WorldNode.pre_environment_step(dt, priority=...)
  -> World.step()
  -> WorldNode.post_environment_step(dt, priority=...)
  -> WorldNode.get_observation()
  -> WorldNode.get_reward() / get_termination() / get_truncation() / get_info() / render()

Implementing a node

Subclasses should:

  1. Set name, world, and any relevant spaces / signal flags / render attributes (render_mode, supported_render_modes) in __init__.
  2. Override the lifecycle methods they need (reset, pre_environment_step, etc.).
  3. Populate the corresponding priority sets for every lifecycle method they override. Callers check these sets before dispatching a method call — a node will only have a given method called for priorities present in its corresponding set. A node with an empty priority set for a given method will never have that method called.

Example::

   class MySensor(WorldNode[...]):
       after_reset_priorities = {0}
       post_environment_step_priorities = {0}

       def post_environment_step(self, dt, *, priority=0):
           ...  # read sensor data

Multiple priorities (e.g. {0, 1}) allow the same method to be invoked at different stages; the caller iterates priorities in the desired order.

Space lifecycle contract

action_space, observation_space, and context_space follow a two-phase lifecycle:

Phase 1 — construction (__init__): Nodes should expose a preliminary space if the dimensionality is known before the scene is built. For example, an EEF controller always produces a 6-D action regardless of the number of robot DOFs, so the space can be set to a BoxSpace with [-inf, inf] bounds immediately. If the dimensionality is genuinely unknown (e.g. a joint-position controller whose DOF count comes from the compiled scene), the space may be None at this stage.

Phase 2 — post-build (end of after_reload): By the time the lowest-priority after_reload has returned, every node must have set its final, tightly-bounded spaces. CombinedWorldNode calls _refresh_spaces() at this point so that the aggregated spaces seen by the environment composer reflect the true, post-build state.


Priority conventions (in reload / reset / after_reload / after_reset): Higher priority runs first (sorted(..., reverse=True)).

  • [100, 200) -> Floor / surrounding environment setup (e.g. table, ground plane)
  • [50, 100) -> Robot setup (add URDF/MJCF entity, init controller, apply rest pose)
  • [0, 50) -> Object / prop setup; general post-step observation updates
  • (-, 0) -> Sensors / renderers that must run after all entities are settled (e.g. cameras: -50)

name instance-attribute

name: str

control_timestep class-attribute instance-attribute

control_timestep: Optional[float] = None

update_timestep class-attribute instance-attribute

update_timestep: Optional[float] = None

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

observation_space class-attribute instance-attribute

observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None

action_space class-attribute instance-attribute

action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None

has_reward class-attribute instance-attribute

has_reward: bool = False

has_termination_signal class-attribute instance-attribute

has_termination_signal: bool = False

has_truncation_signal class-attribute instance-attribute

has_truncation_signal: bool = False

supported_render_modes class-attribute instance-attribute

supported_render_modes: Sequence[str] = ()

render_mode class-attribute instance-attribute

render_mode: Optional[str] = None

world class-attribute instance-attribute

world: Optional[World[BArrayType, BDeviceType, BDtypeType, BRNGType]] = None

reset_priorities class-attribute instance-attribute

reset_priorities: Set[int] = set()

reload_priorities class-attribute instance-attribute

reload_priorities: Set[int] = set()

after_reset_priorities class-attribute instance-attribute

after_reset_priorities: Set[int] = set()

after_reload_priorities class-attribute instance-attribute

after_reload_priorities: Set[int] = set()

pre_environment_step_priorities class-attribute instance-attribute

pre_environment_step_priorities: Set[int] = set()

post_environment_step_priorities class-attribute instance-attribute

post_environment_step_priorities: Set[int] = set()

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

Backend provided by the attached world.

device property

device: Optional[BDeviceType]

Device provided by the attached world.

can_render property

can_render: bool

Whether the node currently exposes a render mode.

effective_update_timestep property

effective_update_timestep: Optional[float]

Resolved update period, defaulting to control_timestep when unset.

unwrapped property

unwrapped: WorldNode

prev_wrapper_layer property

prev_wrapper_layer: Optional[WorldNode]

pre_environment_step

pre_environment_step(dt: Union[float, BArrayType], *, priority: int = 0) -> None

This method is called before the environment step Args: dt (float/BArrayType): The time elapsed between the last world step and the current step (NOT the current step and next step).

get_context

get_context() -> Optional[ContextType]

Get the current context from the node. If the context space is None, this method should not be called.

get_observation

get_observation() -> ObsType

Get the current observation from the sensor. If the observation space is None, this method should not be called.

get_reward

get_reward() -> Union[float, BArrayType]

Get the current reward from the environment. If has_reward is False, this method should not be called.

get_termination

get_termination() -> Union[bool, BArrayType]

Get the current termination signal from the environment. If has_termination_signal is False, this method should not be called.

get_truncation

get_truncation() -> Union[bool, BArrayType]

Get the current truncation signal from the environment. If has_truncation_signal is False, this method should not be called.

get_info

get_info() -> Optional[Dict[str, Any]]

Get optional auxiliary information with this node.

render

render() -> Union[np.ndarray, BArrayType, Sequence[Union[np.ndarray, BArrayType]], Dict[str, Union[np.ndarray, BArrayType, Sequence[Union[np.ndarray, BArrayType]]]], None]

Render the current state of the node.

If can_render is False, this method should not be called.

set_next_action

set_next_action(action: ActType) -> None

Update the next action to be taken by the node. This method should be called before pre_environment_step call. If this method is not called after a world step or an action of None is given in the call, the node will compute a dummy action to try retain the same state of the robot. Note that if the action space is None, this method should not be called.

post_environment_step

post_environment_step(dt: Union[float, BArrayType], *, priority: int = 0) -> None

This method is called after the environment step to update the sensor's internal state. Args: dt (float/BArrayType): The time elapsed between the last world step and the current step.

reset

reset(*, priority: int = 0, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

This method is called after World.reset(...) has been called. Reset the node and update its internal state.

reload

reload(*, priority: int = 0, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None

Reload the node. By default, this just calls reset with the same parameters. Simulation environments can override this to completely re-read assets and reload the node.

after_reset

after_reset(*, priority: int = 0, mask: Optional[BArrayType] = None) -> None

This method is called after all WorldNodes has been called with reset (e.g. the environment reset is effectively done). Use get_context, get_observation, and get_info to read the post-reset state.

after_reload

after_reload(*, priority: int = 0, mask: Optional[BArrayType] = None) -> None

This method is called after the world has been rebuilt following a reload.

At this point: - All nodes have added their entities during the reload phase - The world scene has been built (simulation is ready) - Nodes can now cache references to entities, geoms, etc.

By default, this calls after_reset(). Override if you need specific post-reload initialization that differs from post-reset.

Use get_context, get_observation, and get_info to read the state.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[WorldNode]

Return a node by key path.

  • str keys are treated as a single direct-child key.
  • Empty key sequence returns self.
  • Vanilla nodes have no children and return None for non-empty lookup.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[WorldNode], bool]) -> list[WorldNode]

Return nodes in this subtree matching fn.

Vanilla nodes only evaluate self.

get_nodes_by_type

get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]

Return nodes in this subtree that are instances of node_type.

close

close() -> None

Release any node-owned resources.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

FuncWorld

Bases: ABC, Generic[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]

Functional counterpart to :class:World with explicit state passing.

backend instance-attribute

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device instance-attribute

device: Optional[BDeviceType]

The world timestep in seconds. If None, the world is asynchronous (real-time)

world_timestep instance-attribute

world_timestep: Optional[float]

The world's physical timestep in seconds. There might be multiple world sub-steps inside a world step. If none this means it is not known

world_subtimestep class-attribute instance-attribute

world_subtimestep: Optional[float] = None

The number of parallel environments in this world

batch_size class-attribute instance-attribute

batch_size: Optional[int] = None

initial abstractmethod

initial(*, seed: Optional[int] = None, **kwargs) -> WorldStateT

Construct the initial world state.

step abstractmethod

step(state: WorldStateT) -> Tuple[WorldStateT, Union[float, BArrayType]]

Advance state by one world step.

reset abstractmethod

reset(state: WorldStateT, *, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> WorldStateT

Perform reset on the selected environments with the given mask Note that the state input and output should be with the full batch size

reload

reload(state: WorldStateT, *, seed: Optional[int] = None, **kwargs) -> WorldStateT

Prepare the world for reloading. This is called BEFORE nodes add entities.

Use this hook to: - Clear old scene data - Load asset definitions - Prepare for entities to be added by nodes

DO NOT compile the scene here. Compilation happens in after_reload() after all nodes have had a chance to add their entities.

By default, this just calls initial with the same parameters. Simulation environments should override this to prepare for entity addition.

The reload flow is
  1. world.reload() # World prepares (clear, load assets)
  2. node.reload() # Nodes add entities to scene
  3. world.after_reload() # World compiles scene with all entities
  4. node.after_reload() # Nodes cache references

after_reset

after_reset(state: WorldStateT, *, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> WorldStateT

Called after reset() completes.

Use this hook to perform post-reset initialization. Returns the (possibly modified) world state.

after_reload

after_reload(state: WorldStateT, *, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> WorldStateT

Compile the world scene AFTER all nodes have added their entities.

This is where the actual scene compilation should happen, since nodes have already added their entities during the reload() phase.

Use this hook to: - Compile the simulation scene - Cache references to entities - Perform validation that requires the compiled scene - Initialize state that depends on the final scene configuration

Returns the (possibly modified) world state.

close

close(state: WorldStateT) -> None

Release any resources referenced by state.

is_control_timestep_compatible

is_control_timestep_compatible(control_timestep: Optional[float]) -> bool

Check whether a node control period aligns with world_timestep.

FuncWorldNode

Bases: ABC, Generic[WorldStateT, NodeStateT, ContextType, ObsType, ActType, BArrayType, BDeviceType, BDtypeType, BRNGType]

A functional (stateless) node that manages one aspect of a world.

Unlike WorldNode, every method receives and returns explicit state objects (world_state, node_state) instead of mutating internal attributes.

initial creates the node state for the first time (the environment has not been set up yet). reset is called on subsequent episode boundaries when the node state already exists.

Lifecycle — initial flow (environment not yet created)::

FuncWorldNode.initial(world_state, priority=...)
  -> returns (world_state, node_state)

Lifecycle — reset flow::

FuncWorld.reset()
  -> FuncWorldNode.reset(world_state, node_state, priority=...)
  -> FuncWorldNode.after_reset(world_state, node_state, priority=...)
  -> FuncWorldNode.get_context(...) / get_observation(...) / get_info(...) / render(...)

Lifecycle — reload flow::

FuncWorld.reload()          # World prepares (clear, load assets)
  -> FuncWorldNode.reload(world_state, priority=...)     # Nodes add entities
FuncWorld.after_reload()    # World compiles scene
  -> FuncWorldNode.after_reload(world_state, node_state, priority=...)  # Nodes cache refs
  -> FuncWorldNode.get_context(...) / get_observation(...) / get_info(...) / render(...)

reload re-generates the simulation environment (e.g. re-reading assets, rebuilding the scene). This is typically much more expensive than reset and should only be called when the environment configuration has changed. By default reload delegates to initial.

Lifecycle — step flow::

FuncWorldNode.set_next_action(world_state, node_state, action)
  -> FuncWorldNode.pre_environment_step(world_state, node_state, dt, priority=...)
  -> FuncWorld.step()
  -> FuncWorldNode.post_environment_step(world_state, node_state, dt, priority=...)
  -> FuncWorldNode.get_observation(world_state, node_state)
  -> FuncWorldNode.get_reward(...) / get_termination(...) / get_truncation(...) / get_info(...) / render(...)

Implementing a node

Subclasses should:

  1. Set name, world, and any relevant spaces / signal flags / render attributes (render_mode, supported_render_modes) in __init__ (or as class-level attributes).
  2. Implement initial and reset (both abstract), plus any other lifecycle methods the node needs (pre_environment_step, post_environment_step, etc.).
  3. Populate the corresponding priority sets for every lifecycle method they implement. Callers check these sets before dispatching a method call — a node will only have a given method called for priorities present in its corresponding set. A node with an empty priority set for a given method will never have that method called.

Example::

   class MySensor(FuncWorldNode[...]):
       initial_priorities = {0}
       reset_priorities = {0}
       after_reset_priorities = {0}
       post_environment_step_priorities = {0}

       def initial(self, world_state, *, priority=0, seed=None, **kwargs):
           ...
       def reset(self, world_state, node_state, *, priority=0, seed=None, mask=None, **kwargs):
           ...

Multiple priorities (e.g. {0, 1}) allow the same method to be invoked at different stages; the caller iterates priorities in the desired order.

Space contract

action_space, observation_space, and context_space are static interface metadata. They must be fully declared in __init__ and must not change after construction. CombinedFuncWorldNode reads them once at construction time; there is no refresh mechanism.

Unlike stateful WorldNode subclasses (which may need to build a simulation scene before knowing DOF counts), functional nodes are expected to derive their spaces purely from their configuration parameters.

name instance-attribute

name: str

control_timestep class-attribute instance-attribute

control_timestep: Optional[float] = None

update_timestep class-attribute instance-attribute

update_timestep: Optional[float] = None

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

observation_space class-attribute instance-attribute

observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None

action_space class-attribute instance-attribute

action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None

has_reward class-attribute instance-attribute

has_reward: bool = False

has_termination_signal class-attribute instance-attribute

has_termination_signal: bool = False

has_truncation_signal class-attribute instance-attribute

has_truncation_signal: bool = False

supported_render_modes class-attribute instance-attribute

supported_render_modes: Sequence[str] = ()

render_mode class-attribute instance-attribute

render_mode: Optional[str] = None

world class-attribute instance-attribute

world: Optional[FuncWorld[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None

initial_priorities class-attribute instance-attribute

initial_priorities: Set[int] = set()

reset_priorities class-attribute instance-attribute

reset_priorities: Set[int] = set()

reload_priorities class-attribute instance-attribute

reload_priorities: Set[int] = set()

after_reset_priorities class-attribute instance-attribute

after_reset_priorities: Set[int] = set()

after_reload_priorities class-attribute instance-attribute

after_reload_priorities: Set[int] = set()

pre_environment_step_priorities class-attribute instance-attribute

pre_environment_step_priorities: Set[int] = set()

post_environment_step_priorities class-attribute instance-attribute

post_environment_step_priorities: Set[int] = set()

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

Backend provided by the attached world.

device property

device: Optional[BDeviceType]

Device provided by the attached world.

can_render property

can_render: bool

Whether the node currently exposes a render mode.

effective_update_timestep property

effective_update_timestep: Optional[float]

Resolved update period, defaulting to control_timestep when unset.

unwrapped property

unwrapped: FuncWorldNode

prev_wrapper_layer property

prev_wrapper_layer: Optional[FuncWorldNode]

initial abstractmethod

initial(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, **kwargs) -> Tuple[WorldStateT, NodeStateT]

Create the node state for a freshly initialized world.

reload

reload(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, **kwargs) -> Tuple[WorldStateT, NodeStateT]

Reload the node. By default, this just calls initial with the same parameters. Simulation environments can override this to completely re-read assets and reload the node.

reset abstractmethod

reset(world_state: WorldStateT, node_state: NodeStateT, *, priority: int = 0, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> Tuple[WorldStateT, NodeStateT]

This method is called after FuncWorld.reset(...) has been called.

after_reset

after_reset(world_state: WorldStateT, node_state: NodeStateT, *, priority: int = 0, mask: Optional[BArrayType] = None) -> Tuple[WorldStateT, NodeStateT]

This method is called after all FuncWorldNodes has been called with reset (e.g. the environment reset is effectively done). Use get_context, get_observation, and get_info to read the post-reset state.

after_reload

after_reload(world_state: WorldStateT, node_state: NodeStateT, *, priority: int = 0, mask: Optional[BArrayType] = None) -> Tuple[WorldStateT, NodeStateT]

This method is called after the world has been rebuilt following a reload.

At this point: - All nodes have added their entities during the reload phase - The world scene has been built (simulation is ready) - Nodes can now cache references to entities, geoms, etc.

By default, this calls after_reset(). Override if you need specific post-reload initialization that differs from post-reset.

Use get_context, get_observation, and get_info to read the state.

pre_environment_step

pre_environment_step(world_state: WorldStateT, node_state: NodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, NodeStateT]

This method is called before each environment step. Args: world_state (WorldStateT): The current state of the world. node_state (NodeStateT): The current state of the node. dt (Union[float, BArrayType]): The time delta since the last step.

get_context

get_context(world_state: WorldStateT, node_state: NodeStateT) -> Optional[ContextType]

Get the current context from the node. If the context space is None, this method should not be called.

get_observation

get_observation(world_state: WorldStateT, node_state: NodeStateT) -> ObsType

Get the current observation from the sensor. If the observation space is None, this method should not be called.

get_reward

get_reward(world_state: WorldStateT, node_state: NodeStateT) -> Union[float, BArrayType]

Get the current reward from the environment. If the reward space is None, this method should not be called.

get_termination

get_termination(world_state: WorldStateT, node_state: NodeStateT) -> Union[bool, BArrayType]

Get the current termination status from the environment. If the termination space is None, this method should not be called.

get_truncation

get_truncation(world_state: WorldStateT, node_state: NodeStateT) -> Union[bool, BArrayType]

Get the current truncation status from the environment. If the truncation space is None, this method should not be called.

get_info

get_info(world_state: WorldStateT, node_state: NodeStateT) -> Optional[Dict[str, Any]]

Get the current info from the environment.

render

render(world_state: WorldStateT, node_state: NodeStateT) -> Union[np.ndarray, BArrayType, Sequence[Union[np.ndarray, BArrayType]], Dict[str, Union[np.ndarray, BArrayType, Sequence[Union[np.ndarray, BArrayType]]]], None]

Render the current state of the node.

If can_render is False, this method should not be called.

set_next_action

set_next_action(world_state: WorldStateT, node_state: NodeStateT, action: ActType) -> Tuple[WorldStateT, NodeStateT]

Update the next action to be taken by the node. This method should be called before pre_environment_step call. If this method is not called after a world step or an action of None is given in the call, the node will compute a dummy action to try retain the same state of the robot. Note that if the action space is None, this method should not be called.

post_environment_step

post_environment_step(world_state: WorldStateT, node_state: NodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, NodeStateT]

This method is called after the environment step to update the sensor's internal state. Args: world_state (WorldStateT): The current state of the world. node_state (NodeStateT): The current state of the node. dt (Union[float, BArrayType]): The time delta since the last step.

close

close(world_state: WorldStateT, node_state: NodeStateT) -> WorldStateT

Release node-owned resources and return the final world state.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[FuncWorldNode]

Return a node by key path.

  • str keys are treated as a single direct-child key.
  • Empty key sequence returns self.
  • Vanilla nodes have no children and return None for non-empty lookup.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[FuncWorldNode], bool]) -> list[FuncWorldNode]

Return nodes in this subtree matching fn.

Vanilla nodes only evaluate self.

get_nodes_by_type

get_nodes_by_type(node_type: Type[FuncWorldNode]) -> list[FuncWorldNode]

Return nodes in this subtree that are instances of node_type.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

WorldEnv

WorldEnv(world: World[BArrayType, BDeviceType, BDtypeType, BRNGType], node_or_nodes: Union[WorldNode[ContextType, ObsType, ActType, BArrayType, BDeviceType, BDtypeType, BRNGType], Iterable[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]], *, render_mode: Optional[str] = 'auto')

Bases: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Composes an Env from a World and one or more WorldNode instances.

If multiple nodes are passed, they are automatically wrapped in a CombinedWorldNode.

Bind a world and node tree into one stateful environment facade.

metadata class-attribute instance-attribute

metadata: dict[str, Any] = {'render_modes': []}

unwrapped property

unwrapped: Env

prev_wrapper_layer property

prev_wrapper_layer: Optional[Env]

node instance-attribute

node = node_or_nodes

world instance-attribute

world = world

rng instance-attribute

rng = random_number_generator(device=device)

observation_space property

observation_space

action_space property

action_space

context_space property

context_space

backend property

backend

device property

device

batch_size property

batch_size

render_mode property

render_mode: Optional[str]

render_fps property

render_fps: Optional[int]

sample_space

sample_space(space: Space) -> Any

Sample from space using and updating self.rng.

sample_action

sample_action() -> ActType

Sample one action from action_space.

sample_observation

sample_observation() -> ObsType

Sample one observation from observation_space.

sample_context

sample_context() -> Optional[ContextType]

Sample one context value if context_space is defined.

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Look up a node by name or nested path.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], bool]) -> list[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Collect nodes in the tree that satisfy fn.

get_nodes_by_type

get_nodes_by_type(node_type: Type[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]) -> list[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Collect nodes in the tree that are instances of node_type.

reload

reload(*, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]

Reload the environment - reload scene with entities from nodes.

Flow
  1. world.reload() - World prepares (clear, load assets)
  2. node.reload() - Nodes add entities to scene
  3. world.after_reload() - World compiles scene with all entities
  4. node.after_reload() - Nodes cache entity references

reset

reset(*, mask: Optional[BArrayType] = None, seed: Optional[int] = None, reload: bool = False, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]

Reset the composed environment, optionally forcing a full reload.

step

step(action: ActType) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

Run one control step through the node/world composition.

render

render() -> RenderFrame | Sequence[RenderFrame] | None

Render via the root node when rendering is enabled.

close

close()

Close the node tree and the underlying world.

FuncWorldEnv

FuncWorldEnv(world: FuncWorld[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType], node_or_nodes: Union[FuncWorldNode[WorldStateT, NodeStateT, ContextType, ObsType, ActType, BArrayType, BDeviceType, BDtypeType, BRNGType], Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]], *, render_mode: Optional[str] = 'auto')

Bases: FuncEnv[WorldFuncEnvState, None, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]

Composes a FuncEnv from a FuncWorld and one or more FuncWorldNode instances.

If multiple nodes are passed, they are automatically wrapped in a CombinedFuncWorldNode.

Bind a functional world and node tree into one FuncEnv facade.

metadata class-attribute instance-attribute

metadata: Dict[str, Any] = {'render_modes': []}

unwrapped property

unwrapped: FuncEnv

prev_wrapper_layer property

prev_wrapper_layer: Optional[FuncEnv]

node instance-attribute

node = node_or_nodes

world instance-attribute

world = world

observation_space property

observation_space

action_space property

action_space

context_space property

context_space

backend property

backend

device property

device

batch_size property

batch_size

render_mode property

render_mode: Optional[str]

render_fps property

render_fps: Optional[int]

update_observation_post_reset

update_observation_post_reset(old_obs: ObsType, newobs_masked: ObsType, mask: BArrayType) -> ObsType

Merge masked reset observations back into a full batched observation.

update_context_post_reset

update_context_post_reset(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType

Merge masked reset contexts back into a full batched context.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Look up a node by name or nested path.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], bool]) -> list[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Collect nodes in the tree that satisfy fn.

get_nodes_by_type

get_nodes_by_type(node_type: Type[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]) -> list[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Collect nodes in the tree that are instances of node_type.

reload

reload(state: WorldFuncEnvState, *, seed: Optional[int] = None, **kwargs) -> Tuple[WorldFuncEnvState, ContextType, ObsType, Dict[str, Any]]

Rebuild the world and recreate node state through the reload lifecycle.

initial

initial(*, seed: Optional[int] = None, **kwargs) -> Tuple[WorldFuncEnvState, ContextType, ObsType, Dict[str, Any]]

Create the initial world state and root node state.

reset

reset(state: WorldFuncEnvState, *, seed: Optional[int] = None, mask: Optional[BArrayType] = None, reload: bool = False, **kwargs) -> Tuple[WorldFuncEnvState, ContextType, ObsType, Dict[str, Any]]

Reset the composed functional environment, optionally forcing reload.

step

step(state: WorldFuncEnvState, action: ActType) -> Tuple[WorldFuncEnvState, ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]

Run one control step through the functional node/world composition.

close

close(state: WorldFuncEnvState) -> None

Close the node tree first, then the underlying world.

render_init

render_init(state, *, seed=None, render_mode=None, **kwargs)

Return a no-op render state for node-based functional environments.

render_image

render_image(state, render_state)

Render through the root node when rendering is enabled.

render_close

render_close(state, render_state)

Return the unchanged state because no extra render state is held.

WorldFuncEnvState dataclass

WorldFuncEnvState(world_state: Any, node_state: Any)

State object used by FuncWorldEnv, bundling world and node states.

world_state instance-attribute

world_state: Any

node_state instance-attribute

node_state: Any

CombinedWorldNode

CombinedWorldNode(name: str, nodes: Iterable[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], direct_return: bool = True, render_mode: Optional[str] = 'auto')

Bases: WorldNode[Optional[CombinedDataT], CombinedDataT, CombinedDataT, BArrayType, BDeviceType, BDtypeType, BRNGType], Generic[BArrayType, BDeviceType, BDtypeType, BRNGType]

A WorldNode that combines multiple WorldNodes into one node, using a dictionary to store the data from each node. The observation, reward, termination, truncation, and info are combined from all child nodes. The keys in the dictionary are the names of the child nodes. If there is only one child node that supports value and direct_return is set to True, the value is returned directly instead of a dictionary.

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

observation_space class-attribute instance-attribute

observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None

action_space class-attribute instance-attribute

action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

Backend provided by the attached world.

device property

device: Optional[BDeviceType]

Device provided by the attached world.

can_render property

can_render: bool

Whether the node currently exposes a render mode.

unwrapped property

unwrapped: WorldNode

prev_wrapper_layer property

prev_wrapper_layer: Optional[WorldNode]

supported_render_modes class-attribute instance-attribute

supported_render_modes = ('dict', 'auto')

nodes instance-attribute

nodes = nodes

has_reward instance-attribute

has_reward = any((has_reward) for node in nodes)

has_termination_signal instance-attribute

has_termination_signal = any((has_termination_signal) for node in nodes)

has_truncation_signal instance-attribute

has_truncation_signal = any((has_truncation_signal) for node in nodes)

name instance-attribute

name = name

direct_return instance-attribute

direct_return = direct_return

render_mode instance-attribute

render_mode = render_mode

world property

world: World[BArrayType, BDeviceType, BDtypeType, BRNGType]

control_timestep property

control_timestep: Optional[float]

update_timestep property

update_timestep: Optional[float]

effective_update_timestep property

effective_update_timestep: Optional[float]

reset_priorities property

reset_priorities: Set[int]

reload_priorities property

reload_priorities: Set[int]

after_reset_priorities property

after_reset_priorities: Set[int]

after_reload_priorities property

after_reload_priorities: Set[int]

pre_environment_step_priorities property

pre_environment_step_priorities: Set[int]

post_environment_step_priorities property

post_environment_step_priorities: Set[int]

get_nodes_by_type

get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]

Return nodes in this subtree that are instances of node_type.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

aggregate_spaces staticmethod

aggregate_spaces(spaces: Dict[str, Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]], direct_return: bool = True) -> Tuple[Optional[str], Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]]

Aggregate child spaces into either a passthrough or a DictSpace.

aggregate_data staticmethod

aggregate_data(data: Dict[str, Any], direct_return: bool = True) -> Optional[Union[Dict[str, Any], Any]]

Aggregate child outputs using the same direct-return convention as spaces.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[WorldNode]

Resolve a child node by dotted-path-like key traversal.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[WorldNode], bool]) -> list[WorldNode]

Return every node in the subtree that satisfies fn.

pre_environment_step

pre_environment_step(dt, *, priority: int = 0)

Dispatch pre-step callbacks to child nodes at the matching frequency.

get_context

get_context()

get_observation

get_observation()

get_reward

get_reward()

get_termination

get_termination()

get_truncation

get_truncation()

get_info

get_info() -> Optional[Dict[str, Any]]

render

render()

set_next_action

set_next_action(action)

Route combined actions to child nodes, respecting per-node control rates.

post_environment_step

post_environment_step(dt, *, priority: int = 0)

Dispatch post-step callbacks to child nodes at the matching frequency.

reset

reset(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})

Forward reset calls to children that participate at priority.

reload

reload(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})

Forward reload calls to children that participate at priority.

after_reset

after_reset(*, priority: int = 0, mask=None)

Forward post-reset hooks and reset internal routing counters.

after_reload

after_reload(*, priority: int = 0, mask=None)

Call after_reload on child nodes. Similar to after_reset but for reload flow.

close

close()

Close every child node.

FlatCombinedWorldNode

FlatCombinedWorldNode(name: str, nodes: Iterable[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], render_mode: Optional[str] = 'auto')

Bases: CombinedWorldNode[BArrayType, BDeviceType, BDtypeType, BRNGType]

A WorldNode that combines multiple WorldNodes and flattens their data.

Unlike CombinedWorldNode which stores data as {node_name: {key: value}}, FlatCombinedWorldNode merges dictionaries directly as {key1: value1, key2: value2}.

This requires: - All nodes with observation/action/context spaces must use DictSpace - Keys across all nodes must be unique (no overlaps)

The node names are only used for identification, not for nesting data.

Initialize a FlatCombinedWorldNode.

Parameters:

Name Type Description Default
name str

Name of this combined node

required
nodes Iterable[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Iterable of nodes to combine

required
render_mode Optional[str]

Render mode ('dict', 'auto', or specific mode)

'auto'

Raises:

Type Description
ValueError

If node spaces have overlapping keys or non-DictSpace types

name instance-attribute

name = name

control_timestep property

control_timestep: Optional[float]

update_timestep property

update_timestep: Optional[float]

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

observation_space class-attribute instance-attribute

observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None

action_space class-attribute instance-attribute

action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None

has_reward instance-attribute

has_reward = any((has_reward) for node in nodes)

has_termination_signal instance-attribute

has_termination_signal = any((has_termination_signal) for node in nodes)

has_truncation_signal instance-attribute

has_truncation_signal = any((has_truncation_signal) for node in nodes)

supported_render_modes class-attribute instance-attribute

supported_render_modes = ('dict', 'auto')

render_mode instance-attribute

render_mode = render_mode

world property

world: World[BArrayType, BDeviceType, BDtypeType, BRNGType]

reset_priorities property

reset_priorities: Set[int]

reload_priorities property

reload_priorities: Set[int]

after_reset_priorities property

after_reset_priorities: Set[int]

after_reload_priorities property

after_reload_priorities: Set[int]

pre_environment_step_priorities property

pre_environment_step_priorities: Set[int]

post_environment_step_priorities property

post_environment_step_priorities: Set[int]

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

Backend provided by the attached world.

device property

device: Optional[BDeviceType]

Device provided by the attached world.

can_render property

can_render: bool

Whether the node currently exposes a render mode.

effective_update_timestep property

effective_update_timestep: Optional[float]

unwrapped property

unwrapped: WorldNode

prev_wrapper_layer property

prev_wrapper_layer: Optional[WorldNode]

nodes instance-attribute

nodes = nodes

direct_return instance-attribute

direct_return = direct_return

pre_environment_step

pre_environment_step(dt, *, priority: int = 0)

Dispatch pre-step callbacks to child nodes at the matching frequency.

get_reward

get_reward()

get_termination

get_termination()

get_truncation

get_truncation()

render

render()

post_environment_step

post_environment_step(dt, *, priority: int = 0)

Dispatch post-step callbacks to child nodes at the matching frequency.

reset

reset(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})

Forward reset calls to children that participate at priority.

reload

reload(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})

Forward reload calls to children that participate at priority.

after_reset

after_reset(*, priority: int = 0, mask=None)

Forward post-reset hooks and reset internal routing counters.

after_reload

after_reload(*, priority: int = 0, mask=None)

Call after_reload on child nodes. Similar to after_reset but for reload flow.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[WorldNode]

Resolve a child node by dotted-path-like key traversal.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[WorldNode], bool]) -> list[WorldNode]

Return every node in the subtree that satisfies fn.

get_nodes_by_type

get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]

Return nodes in this subtree that are instances of node_type.

close

close()

Close every child node.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

aggregate_spaces staticmethod

aggregate_spaces(spaces: Dict[str, Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]], direct_return: bool = True) -> Tuple[Optional[str], Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]]

Aggregate child spaces into either a passthrough or a DictSpace.

aggregate_data staticmethod

aggregate_data(data: Dict[str, Any], direct_return: bool = True) -> Optional[Union[Dict[str, Any], Any]]

Aggregate child outputs using the same direct-return convention as spaces.

get_context

get_context() -> Optional[CombinedDataT]

Get context by flattening all node contexts into one dictionary.

get_observation

get_observation() -> CombinedDataT

Get observation by flattening all node observations into one dictionary.

get_info

get_info() -> Optional[Dict[str, Any]]

Get info by merging all node info dictionaries.

set_next_action

set_next_action(action: CombinedDataT) -> None

Set action by routing keys to appropriate nodes.

CombinedFuncWorldNode

CombinedFuncWorldNode(name: str, nodes: Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], direct_return: bool = True, render_mode: Optional[str] = 'auto')

Bases: FuncWorldNode[WorldStateT, CombinedNodeStateT, Optional[CombinedDataT], CombinedDataT, CombinedDataT, BArrayType, BDeviceType, BDtypeType, BRNGType], Generic[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]

A functional counterpart to CombinedWorldNode that composes multiple FuncWorldNodes.

It aggregates spaces (context, observation, action) and runtime data (context, observation, info, reward, termination, truncation) across child nodes. If only one child exposes a given interface and direct_return=True, the value is passed through directly.

context_space class-attribute instance-attribute

context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None

observation_space class-attribute instance-attribute

observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None

action_space class-attribute instance-attribute

action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

Backend provided by the attached world.

device property

device: Optional[BDeviceType]

Device provided by the attached world.

can_render property

can_render: bool

Whether the node currently exposes a render mode.

unwrapped property

unwrapped: FuncWorldNode

prev_wrapper_layer property

prev_wrapper_layer: Optional[FuncWorldNode]

supported_render_modes class-attribute instance-attribute

supported_render_modes = ('dict', 'auto')

nodes instance-attribute

nodes = nodes

has_reward instance-attribute

has_reward = any((has_reward) for node in nodes)

has_termination_signal instance-attribute

has_termination_signal = any((has_termination_signal) for node in nodes)

has_truncation_signal instance-attribute

has_truncation_signal = any((has_truncation_signal) for node in nodes)

name instance-attribute

name = name

direct_return instance-attribute

direct_return = direct_return

render_mode instance-attribute

render_mode = render_mode

world property

world

control_timestep property

control_timestep

update_timestep property

update_timestep

effective_update_timestep property

effective_update_timestep

initial_priorities property

initial_priorities: Set[int]

reset_priorities property

reset_priorities: Set[int]

reload_priorities property

reload_priorities: Set[int]

after_reset_priorities property

after_reset_priorities: Set[int]

after_reload_priorities property

after_reload_priorities: Set[int]

pre_environment_step_priorities property

pre_environment_step_priorities: Set[int]

post_environment_step_priorities property

post_environment_step_priorities: Set[int]

get_nodes_by_type

get_nodes_by_type(node_type: Type[FuncWorldNode]) -> list[FuncWorldNode]

Return nodes in this subtree that are instances of node_type.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

aggregate_spaces staticmethod

aggregate_spaces(spaces: Dict[str, Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]], direct_return: bool = True) -> Tuple[Optional[str], Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]]

Aggregate child spaces into either a passthrough or a DictSpace.

aggregate_data staticmethod

aggregate_data(data: Dict[str, Any], direct_return: bool = True) -> Optional[Union[Dict[str, Any], Any]]

Aggregate child outputs using the same direct-return convention as spaces.

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[FuncWorldNode]

Resolve a child node by dotted-path-like key traversal.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[FuncWorldNode], bool]) -> list[FuncWorldNode]

Return every node in the subtree that satisfies fn.

initial

initial(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]

Create child node states for the current initialization priority.

reload

reload(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]

Recreate child node states for the current reload priority.

reset

reset(world_state: WorldStateT, node_state: CombinedNodeStateT, *, priority: int = 0, seed: Optional[int] = None, mask: Optional[BArrayType] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}, **kwargs) -> Tuple[WorldStateT, CombinedNodeStateT]

Forward reset calls to children that participate at priority.

after_reset

after_reset(world_state: WorldStateT, node_state: CombinedNodeStateT, *, priority: int = 0, mask: Optional[BArrayType] = None) -> Tuple[WorldStateT, CombinedNodeStateT]

Forward post-reset hooks and reset internal routing counters.

after_reload

after_reload(world_state: WorldStateT, node_state: CombinedNodeStateT, *, priority: int = 0, mask: Optional[BArrayType] = None) -> Tuple[WorldStateT, CombinedNodeStateT]

Call after_reload on child nodes. Similar to after_reset but for reload flow.

pre_environment_step

pre_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]

set_next_action

set_next_action(world_state: WorldStateT, node_state: CombinedNodeStateT, action: CombinedDataT) -> Tuple[WorldStateT, CombinedNodeStateT]

post_environment_step

post_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]

close

close(world_state: WorldStateT, node_state: CombinedNodeStateT) -> WorldStateT

get_context

get_context(world_state: WorldStateT, node_state: CombinedNodeStateT) -> CombinedDataT

get_observation

get_observation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> CombinedDataT

get_reward

get_reward(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[float, BArrayType]

get_termination

get_termination(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]

get_truncation

get_truncation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]

get_info

get_info(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Optional[Dict[str, Any]]

render

render(world_state, node_state)

FlatCombinedFuncWorldNode

FlatCombinedFuncWorldNode(name: str, nodes: Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], render_mode: Optional[str] = 'auto')

Bases: CombinedFuncWorldNode[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]

A FuncWorldNode that combines multiple FuncWorldNodes and flattens their data.

Unlike CombinedFuncWorldNode which stores data as {node_name: {key: value}}, FlatCombinedFuncWorldNode merges dictionaries directly as {key1: value1, key2: value2}.

This requires: - All nodes with observation/action/context spaces must use DictSpace - Keys across all nodes must be unique (no overlaps)

The node names are only used for identification, not for nesting data.

Initialize a FlatCombinedFuncWorldNode.

Parameters:

Name Type Description Default
name str

Name of this combined node

required
nodes Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]

Iterable of nodes to combine

required
render_mode Optional[str]

Render mode ('dict', 'auto', or specific mode)

'auto'

Raises:

Type Description
ValueError

If node spaces have overlapping keys or non-DictSpace types

name instance-attribute

name = name

control_timestep property

control_timestep

update_timestep property

update_timestep

has_reward instance-attribute

has_reward = any((has_reward) for node in nodes)

has_termination_signal instance-attribute

has_termination_signal = any((has_termination_signal) for node in nodes)

has_truncation_signal instance-attribute

has_truncation_signal = any((has_truncation_signal) for node in nodes)

supported_render_modes class-attribute instance-attribute

supported_render_modes = ('dict', 'auto')

render_mode instance-attribute

render_mode = render_mode

world property

world

initial_priorities property

initial_priorities: Set[int]

reset_priorities property

reset_priorities: Set[int]

reload_priorities property

reload_priorities: Set[int]

after_reset_priorities property

after_reset_priorities: Set[int]

after_reload_priorities property

after_reload_priorities: Set[int]

pre_environment_step_priorities property

pre_environment_step_priorities: Set[int]

post_environment_step_priorities property

post_environment_step_priorities: Set[int]

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

Backend provided by the attached world.

device property

device: Optional[BDeviceType]

Device provided by the attached world.

can_render property

can_render: bool

Whether the node currently exposes a render mode.

effective_update_timestep property

effective_update_timestep

unwrapped property

unwrapped: FuncWorldNode

prev_wrapper_layer property

prev_wrapper_layer: Optional[FuncWorldNode]

nodes instance-attribute

nodes = nodes

direct_return instance-attribute

direct_return = direct_return

context_space instance-attribute

context_space = _flatten_spaces([(context_space) for node in (nodes) if context_space is not None], ignore_duplicate_keys=False)

observation_space instance-attribute

observation_space = _flatten_spaces([(observation_space) for node in (nodes) if observation_space is not None], ignore_duplicate_keys=False)

action_space instance-attribute

action_space = _flatten_spaces([(action_space) for node in (nodes) if action_space is not None], ignore_duplicate_keys=True)

initial

initial(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]

Create child node states for the current initialization priority.

reload

reload(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]

Recreate child node states for the current reload priority.

reset

reset(world_state: WorldStateT, node_state: CombinedNodeStateT, *, priority: int = 0, seed: Optional[int] = None, mask: Optional[BArrayType] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}, **kwargs) -> Tuple[WorldStateT, CombinedNodeStateT]

Forward reset calls to children that participate at priority.

after_reset

after_reset(world_state: WorldStateT, node_state: CombinedNodeStateT, *, priority: int = 0, mask: Optional[BArrayType] = None) -> Tuple[WorldStateT, CombinedNodeStateT]

Forward post-reset hooks and reset internal routing counters.

after_reload

after_reload(world_state: WorldStateT, node_state: CombinedNodeStateT, *, priority: int = 0, mask: Optional[BArrayType] = None) -> Tuple[WorldStateT, CombinedNodeStateT]

Call after_reload on child nodes. Similar to after_reset but for reload flow.

pre_environment_step

pre_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]

get_reward

get_reward(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[float, BArrayType]

get_termination

get_termination(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]

get_truncation

get_truncation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]

render

render(world_state, node_state)

post_environment_step

post_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]

close

close(world_state: WorldStateT, node_state: CombinedNodeStateT) -> WorldStateT

get_node

get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[FuncWorldNode]

Resolve a child node by dotted-path-like key traversal.

get_nodes_by_fn

get_nodes_by_fn(fn: Callable[[FuncWorldNode], bool]) -> list[FuncWorldNode]

Return every node in the subtree that satisfies fn.

get_nodes_by_type

get_nodes_by_type(node_type: Type[FuncWorldNode]) -> list[FuncWorldNode]

Return nodes in this subtree that are instances of node_type.

has_wrapper_attr

has_wrapper_attr(name: str) -> bool

Checks if the attribute name exists in the environment.

get_wrapper_attr

get_wrapper_attr(name: str) -> Any

Gets the attribute name from the environment.

set_wrapper_attr

set_wrapper_attr(name: str, value: Any)

Sets the attribute name on the environment with value.

aggregate_spaces staticmethod

aggregate_spaces(spaces: Dict[str, Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]], direct_return: bool = True) -> Tuple[Optional[str], Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]]

Aggregate child spaces into either a passthrough or a DictSpace.

aggregate_data staticmethod

aggregate_data(data: Dict[str, Any], direct_return: bool = True) -> Optional[Union[Dict[str, Any], Any]]

Aggregate child outputs using the same direct-return convention as spaces.

get_context

get_context(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Optional[CombinedDataT]

Get context by flattening all node contexts into one dictionary.

get_observation

get_observation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> CombinedDataT

Get observation by flattening all node observations into one dictionary.

get_info

get_info(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Optional[Dict[str, Any]]

Get info by merging all node info dictionaries.

set_next_action

set_next_action(world_state: WorldStateT, node_state: CombinedNodeStateT, action: CombinedDataT) -> tuple[WorldStateT, CombinedNodeStateT]

Set action by routing keys to appropriate nodes.

DataTransformation

Bases: ABC

Abstract transformation between two UniEnv data spaces.

Transformations are used by wrappers, storages, and batch views to adapt data while keeping enough metadata to derive the target space and, when supported, reconstruct the inverse transform.

has_inverse class-attribute instance-attribute

has_inverse: bool = False

get_target_space_from_source abstractmethod

get_target_space_from_source(source_space: Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]) -> Optional[Space[TargetDataT, BDeviceType, BDtypeType, BRNGType]]

Returns the target space based on the source space (if supported, otherwise throws ValueError). This is useful for transformations that depend on the source space.

transform abstractmethod

transform(source_space: Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT], data: SourceDataT) -> TargetDataT

Apply the transformation to data drawn from source_space.

direction_inverse

direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]

Return the inverse transformation when one exists.

close

close()

serialize abstractmethod

serialize(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Dict[str, Any]

Serialize the transformation to a JSON-compatible dictionary.

Parameters:

Name Type Description Default
source_space Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]]

Optional source space context for transformations whose serialization depends on backend/device information.

None

Returns:

Name Type Description
dict Dict[str, Any]

A dictionary representation of the transformation containing: - Parameters specific to the transformation type. The helper transformation_to_json is responsible for adding the serialized transformation type field.

deserialize_from abstractmethod classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> DataTransformation

Deserialize a transformation from a JSON-compatible dictionary.

Parameters:

Name Type Description Default
json_data Dict[str, Any]

The dictionary containing the transformation data

required
source_space Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]

Optional source space context for transformations whose deserialization depends on backend/device information.

None

Returns:

Name Type Description
DataTransformation DataTransformation

A new instance of the transformation

IdentityTransformation

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space=None)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> IdentityTransformation

RescaleTransformation

RescaleTransformation(new_low: Union[BArrayType, float] = -1.0, new_high: Union[BArrayType, float] = 1.0, new_dtype: Optional[BDtypeType] = None, nan_to: Optional[Union[float, int, BArrayType]] = None)

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

new_low instance-attribute

new_low = new_low

new_high instance-attribute

new_high = new_high

new_dtype instance-attribute

new_dtype = new_dtype

nan_to instance-attribute

nan_to = nan_to

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space=None)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> RescaleTransformation

DictIncludeKeyTransformation

DictIncludeKeyTransformation(enabled_keys: Iterable[str], *, ignore_missing_keys: bool = False, nested_separator: str = '/')

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = False

nested_separator instance-attribute

nested_separator = nested_separator

enabled_keys property writable

enabled_keys: Set[str]

ignore_missing_keys instance-attribute

ignore_missing_keys = ignore_missing_keys

direction_inverse

direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]

Return the inverse transformation when one exists.

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> DictIncludeKeyTransformation

DictExcludeKeyTransformation

DictExcludeKeyTransformation(excluded_keys: Iterable[str], *, ignore_missing_keys: bool = False, nested_separator: str = '/')

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = False

nested_separator instance-attribute

nested_separator = nested_separator

excluded_keys property writable

excluded_keys: Set[str]

ignore_missing_keys instance-attribute

ignore_missing_keys = ignore_missing_keys

direction_inverse

direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]

Return the inverse transformation when one exists.

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> DictExcludeKeyTransformation

BatchifyTransformation

BatchifyTransformation(axis: int = 0)

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

axis instance-attribute

axis = axis

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space=None)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> BatchifyTransformation

UnBatchifyTransformation

UnBatchifyTransformation(axis: int = 0)

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

axis instance-attribute

axis = axis

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space=None)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> UnBatchifyTransformation

DictTransformation

DictTransformation(mapping: Dict[str, DataTransformation], ignore_missing_keys: bool = False, nested_separator: str = '/')

Bases: DataTransformation

mapping instance-attribute

mapping = mapping

ignore_missing_keys instance-attribute

ignore_missing_keys = ignore_missing_keys

nested_separator instance-attribute

nested_separator = nested_separator

has_inverse instance-attribute

has_inverse = all((has_inverse) for transformation in (values()))

get_target_space_from_source

get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType])

transform

transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]

direction_inverse

direction_inverse(source_space=None) -> Optional[DictTransformation]

close

close()

serialize

serialize(source_space: Optional[Space] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> DictTransformation

FlattenDictTransformation

FlattenDictTransformation(nested_separator: str = '/')

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

nested_separator instance-attribute

nested_separator = nested_separator

close

close()

get_target_space_from_source

get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType]) -> DictSpace[BDeviceType, BDtypeType, BRNGType]

transform

transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]

direction_inverse

direction_inverse(source_space: Optional[Space] = None) -> UnflattenDictTransformation

serialize

serialize(source_space: Optional[Space] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> FlattenDictTransformation

UnflattenDictTransformation

UnflattenDictTransformation(nested_separator: str = '/')

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

nested_separator instance-attribute

nested_separator = nested_separator

close

close()

get_target_space_from_source

get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType]) -> DictSpace[BDeviceType, BDtypeType, BRNGType]

transform

transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]

direction_inverse

direction_inverse(source_space: Optional[Space] = None) -> FlattenDictTransformation

serialize

serialize(source_space: Optional[Space] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> UnflattenDictTransformation

ChainedTransformation

ChainedTransformation(transformations: Iterable[DataTransformation])

Bases: DataTransformation

transformations instance-attribute

transformations = list(transformations)

has_inverse instance-attribute

has_inverse = all((has_inverse) for transformation in (transformations))

get_target_space_from_source

get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType])

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space: Optional[Space] = None) -> Optional[ChainedTransformation]

close

close()

serialize

serialize(source_space: Optional[Space] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> ChainedTransformation

CropTransformation

CropTransformation(crop_low: Union[int, float, BArrayType], crop_high: Union[int, float, BArrayType])

Bases: DataTransformation

Initialize Crop Transformation. Args: crop_low: Lower bound for cropping the data. crop_high: Upper bound for cropping the data.

has_inverse class-attribute instance-attribute

has_inverse = True

crop_low instance-attribute

crop_low = crop_low

crop_high instance-attribute

crop_high = crop_high

close

close()

validate_source_space

validate_source_space(source_space: Space[Any, BDeviceType, BDtypeType, BRNGType]) -> None

get_crop_range

get_crop_range(source_space: BoxSpace[Any, BDeviceType, BDtypeType, BRNGType]) -> Tuple[BArrayType, BArrayType]

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space=None)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> CropTransformation

ImageResizeTransformation

ImageResizeTransformation(new_height: int, new_width: int)

Bases: DataTransformation

has_inverse class-attribute instance-attribute

has_inverse = True

new_height instance-attribute

new_height = new_height

new_width instance-attribute

new_width = new_width

close

close()

get_target_space_from_source

get_target_space_from_source(source_space)

transform

transform(source_space, data)

direction_inverse

direction_inverse(source_space=None)

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> ImageResizeTransformation

IterativeTransformation

IterativeTransformation(transformation: DataTransformation, is_leaf_node_fn: Callable[[Space[Any, BDeviceType, BDtypeType, BRNGType]], bool] = default_is_leaf_fn, inv_is_leaf_node_fn: Callable[[Space[Any, BDeviceType, BDtypeType, BRNGType]], bool] = default_is_leaf_fn)

Bases: DataTransformation

transformation instance-attribute

transformation = transformation

is_leaf_node_fn instance-attribute

is_leaf_node_fn = is_leaf_node_fn

inv_is_leaf_node_fn instance-attribute

inv_is_leaf_node_fn = inv_is_leaf_node_fn

has_inverse instance-attribute

has_inverse = has_inverse

get_target_space_from_source

get_target_space_from_source(source_space: Space[Any, BDeviceType, BDtypeType, BRNGType])

transform

transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]

direction_inverse

direction_inverse(source_space=None) -> Optional[IterativeTransformation]

close

close()

serialize

serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]

deserialize_from classmethod

deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> IterativeTransformation

transformation_to_json

transformation_to_json(transformation: DataTransformation, source_space: Optional[Space] = None) -> Dict[str, Any]

Serialize a DataTransformation to a JSON-compatible dictionary.

Parameters:

Name Type Description Default
transformation DataTransformation

The transformation to serialize.

required
source_space Optional[Space]

Optional source space context used by transforms that need backend/device information during serialization.

None

Returns:

Name Type Description
dict Dict[str, Any]

A JSON-compatible representation of the transformation.

json_to_transformation

json_to_transformation(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> DataTransformation

Deserialize a DataTransformation from a JSON-compatible dictionary.

Parameters:

Name Type Description Default
json_data Dict[str, Any]

The dictionary containing the transformation data.

required
source_space Optional[Space]

Optional source space context used by transforms that need backend/device information during deserialization.

None

Returns:

Name Type Description
DataTransformation DataTransformation

The deserialized transformation.