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.
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
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.
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.
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.
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)
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.
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.
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.
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))
render
abstractmethod
¶
render() -> RenderFrame | Sequence[RenderFrame] | None
Render the current environment state using the configured render mode.
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.
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_wait
¶
step_wait() -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]
Collect the pending step results and batch them.
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.
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]]
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_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]]
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.
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]]
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]
reset
¶
reset(*args, mask: Optional[WrapperBArrayT] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]
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.
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]]
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_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.
observation_space
instance-attribute
¶
observation_space: Space[Any, BDeviceType, BDtypeType, BRNGType]
context_space
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
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.
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.
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.
context_space
property
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]]
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]]
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.
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]]
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]]
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.
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.
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
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
- world.reload() # World prepares (clear, load assets)
- node.reload() # Nodes add entities to scene
- world.after_reload() # World compiles scene with all entities
- 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.
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.
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
- world.reload() # World prepares (clear, load assets)
- node.reload() # Nodes add entities to scene
- world.after_reload() # World compiles scene with all entities
- 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.
is_control_timestep_compatible
¶
is_control_timestep_compatible(control_timestep: Optional[float]) -> bool
Check whether a node control period aligns with world_timestep.
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:
- Set
name,world, and any relevant spaces / signal flags / render attributes (render_mode,supported_render_modes) in__init__. - Override the lifecycle methods they need (
reset,pre_environment_step, etc.). - 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)
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
supported_render_modes
class-attribute
instance-attribute
¶
supported_render_modes: Sequence[str] = ()
world
class-attribute
instance-attribute
¶
world: Optional[World[BArrayType, BDeviceType, BDtypeType, BRNGType]] = None
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.
effective_update_timestep
property
¶
effective_update_timestep: Optional[float]
Resolved update period, defaulting to control_timestep when unset.
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.
strkeys are treated as a single direct-child key.- Empty key sequence returns
self. - Vanilla nodes have no children and return
Nonefor 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.
has_wrapper_attr
¶
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in 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.
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
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
- world.reload() # World prepares (clear, load assets)
- node.reload() # Nodes add entities to scene
- world.after_reload() # World compiles scene with all entities
- 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.
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:
- Set
name,world, and any relevant spaces / signal flags / render attributes (render_mode,supported_render_modes) in__init__(or as class-level attributes). - Implement
initialandreset(both abstract), plus any other lifecycle methods the node needs (pre_environment_step,post_environment_step, etc.). - 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.
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
supported_render_modes
class-attribute
instance-attribute
¶
supported_render_modes: Sequence[str] = ()
world
class-attribute
instance-attribute
¶
world: Optional[FuncWorld[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None
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.
effective_update_timestep
property
¶
effective_update_timestep: Optional[float]
Resolved update period, defaulting to control_timestep when unset.
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.
strkeys are treated as a single direct-child key.- Empty key sequence returns
self. - Vanilla nodes have no children and return
Nonefor 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.
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.
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.
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
- world.reload() - World prepares (clear, load assets)
- node.reload() - Nodes add entities to scene
- world.after_reload() - World compiles scene with all entities
- 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.
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.
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.
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)
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.
supported_render_modes
class-attribute
instance-attribute
¶
supported_render_modes = ('dict', 'auto')
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)
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.
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.
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.
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 |
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_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')
backend
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
pre_environment_step
¶
pre_environment_step(dt, *, priority: int = 0)
Dispatch pre-step callbacks to child nodes at the matching frequency.
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.
has_wrapper_attr
¶
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in 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.
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.
supported_render_modes
class-attribute
instance-attribute
¶
supported_render_modes = ('dict', 'auto')
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)
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.
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]
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]]
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 |
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')
backend
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
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]
post_environment_step
¶
post_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]
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.
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.
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.
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 |
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
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
DictIncludeKeyTransformation
¶
DictIncludeKeyTransformation(enabled_keys: Iterable[str], *, ignore_missing_keys: bool = False, nested_separator: str = '/')
Bases: DataTransformation
direction_inverse
¶
direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]
Return the inverse transformation when one exists.
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
direction_inverse
¶
direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]
Return the inverse transformation when one exists.
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
UnBatchifyTransformation
¶
UnBatchifyTransformation(axis: int = 0)
Bases: DataTransformation
DictTransformation
¶
DictTransformation(mapping: Dict[str, DataTransformation], ignore_missing_keys: bool = False, nested_separator: str = '/')
Bases: DataTransformation
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]
deserialize_from
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> DictTransformation
FlattenDictTransformation
¶
FlattenDictTransformation(nested_separator: str = '/')
Bases: DataTransformation
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
deserialize_from
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> FlattenDictTransformation
UnflattenDictTransformation
¶
UnflattenDictTransformation(nested_separator: str = '/')
Bases: DataTransformation
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
deserialize_from
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> UnflattenDictTransformation
ChainedTransformation
¶
ChainedTransformation(transformations: Iterable[DataTransformation])
Bases: DataTransformation
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])
direction_inverse
¶
direction_inverse(source_space: Optional[Space] = None) -> Optional[ChainedTransformation]
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.
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]
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
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
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]
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. |