unienv_interface¶
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.
Batched-environment space invariant¶
The observation_space, action_space, and context_space (when
present) always describe the data that the environment actually returns
and accepts – including any batch dimension.
batch_size is None→ the environment is unbatched. The spaces describe single-instance values (no leading batch axis).batch_size == N(includingN == 1) → the environment is batched with a leading batch dimension of size N. The spaces already include that leading batch dimension. Callers, wrappers, and downstream code must not apply an additionalbatch_space(or equivalent) on top of these spaces – doing so would double-batch and produce incorrect shapes.
This invariant matters whenever code builds source/target spaces from the env (e.g. replay-buffer builders, data transformations, or wrappers that unbatch per-slot data): always use the env-side spaces as-is and only unbatch/slice the data, never re-batch the space.
instance-attribute
¶
observation_space: Space[ObsType, BDeviceType, BDtypeType, BRNGType]
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
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.
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.
abstractmethod
¶
render() -> RenderFrame | Sequence[RenderFrame] | None
Render the current environment state using the configured render mode.
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
Seeding: reset(seed=s) reseeds rng with s and resets
sub-environment i with seed s + i (i is the index in the
full batch, also when only a masked subset is reset), so every
sub-environment gets a distinct, reproducible seed.
Instantiate every environment immediately and batch their interfaces.
instance-attribute
¶
action_space = _batch_env_spaces([env.action_space for env in self.envs], device=first_env.device)
instance-attribute
¶
observation_space = _batch_env_spaces([env.observation_space for env in self.envs], device=first_env.device)
instance-attribute
¶
context_space = _batch_env_spaces([env.context_space for env in self.envs], device=first_env.device)
instance-attribute
¶
rng = self.backend.random.random_number_generator(seed, device=first_env.device)
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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(action: ActType) -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]
Step every worker environment once and concatenate the outputs.
render() -> Sequence[RenderFrame] | None
Collect render outputs from every worker that returns a frame.
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, shared_memory: bool = True, copy: bool = True)
Bases: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]
Run multiple non-batched environments in separate worker processes.
Seeding follows :class:SyncVecEnv: reset(seed=s) resets worker i
with seed s + i. If a worker raises, the exception is re-raised in the
parent as a RuntimeError (chained to the worker's exception) and that
worker process exits.
No environment is built in the parent process: the spaces, backend, device,
metadata and render settings are reported by the workers (spaces travel as
:func:space_to_json payloads, so this works with every start method), and
the constructor raises if a worker fails to build its environment or if the
workers disagree on backend or device. Differing sub-environment spaces are
batched like in :class:SyncVecEnv.
Shared memory (shared_memory=True, the default): the parent
allocates one buffer of shape (num_envs, *leaf_shape) per leaf of the
single-env observation, context and action spaces. Worker i writes its
observation / context into row i and reads its action from row i
(step writes the whole action batch once and sends a payload-free
command), so only reward, termination flags and info travel through
the pipe and nothing is pickled or concatenated. Each group (observation,
context, action) that cannot be shared falls back to the pipe on its own,
with one INFO message on the unienv_interface.env_base.vec_env logger;
:attr:shared_groups lists the groups that are shared. shared_memory=False
sends everything through the pipe. Where the buffers live follows the
worker env:
================ =============================== =================================
worker env buffer parent read
================ =============================== =================================
numpy host shared memory zero-copy numpy view
torch CPU host shared memory zero-copy torch view
torch CUDA CUDA buffer owned by the zero-copy on the GPU
parent, shared via CUDA IPC
jax CPU host shared memory copy=True: one copy into a
fresh jax array; copy=False:
zero-copy via DLPack (falls back
to a copy if XLA refuses)
jax GPU the same CUDA IPC buffer, with zero-copy via DLPack
torch as carrier (DLPack in/out)
jax TPU / other host shared memory (the worker one host-to-device transfer per
devices copies device -> host) read
================ =============================== =================================
Actions use the same buffers in the other direction; every worker hands its env a private copy of its row.
Details of shared mode:
- Only fixed-size spaces can be shared:
BoxSpace/BinarySpaceleaves nested inDictSpace/TupleSpace(other groups use the pipe). Observations and contexts must match their declared space exactly (shape and dtype per leaf); a mismatch raises in the worker and surfaces as the usual chainedRuntimeErrornaming the leaf, nothing is cast silently (fix the env or passshared_memory=False). Actions must have shape(num_envs, *action_leaf_shape); their dtype is cast with same-kind rules (e.g. float64 policy outputs into a float32 space).infoalways travels through the pipe. copy=True(default) returns data that never changes afterwards.copy=Falsereturns the shared buffers themselves (or zero-copy views of them): they are only valid until the nextstep/reset(which overwrite them in place) and must not be written to; copy anything you want to keep, and make sure asynchronous consumers (e.g. JAX computations) have finished before the next call. CUDA buffers are synchronised by the parent before everystep/reset. Masked resets always return freshly gathered rows.- Every start method works. With
fork(the Linux default) CUDA IPC works as long as the parent had not initialised CUDA before the workers were forked (no env is built in the parent, so that is the usual case); otherwise the workers cannot use CUDA at all, which torch reports when they build their env - usespawn/forkserverthen. CUDA IPC is Linux-only; elsewhere CUDA buffers are staged through host shared memory (with a warning). JAX GPU envs need PyTorch as the IPC carrier and a CUDA-enabled jaxlib in the parent (otherwise they use the pipe). A TPU belongs to a single process, so TPU envs cannot be shared: run the workers on the host (or let them copy to the host) and the parent performs one host-to-TPU transfer per read. - Workers set
XLA_PYTHON_CLIENT_PREALLOCATE=falseunless the variable is already set, so that N JAX workers do not each reserve 75% of the GPU. - The parent owns the buffers until :meth:
close, which first shuts the workers down; host shared memory is unlinked onclose(also after a worker crash) and when the env is garbage collected, and shared-memory envs that are still open at interpreter exit are closed then. Views returned withcopy=Falsestay readable afterclose(the mapping is released with the last view). A worker that is killed hard (e.g.SIGKILL) cannot release its CUDA IPC references, in which case torch may warn about it when the parent exits.
Spawn worker processes and batch the spaces they report.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
env_fn
|
Iterable[Callable[[], Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]]]
|
One factory per sub-environment; each builds an unbatched env. |
required |
seed
|
Optional[int]
|
Seed of the vector env's |
None
|
ctx
|
Optional[BaseContext]
|
Multiprocessing context (default: the platform default). |
None
|
daemon
|
bool
|
Whether the worker processes are daemonic. |
True
|
shared_memory
|
bool
|
Transfer observations, contexts and actions through
shared buffers wherever their spaces allow it, falling back to the
pipe per group (see the class docstring); |
True
|
copy
|
bool
|
Only used with shared memory: return copies ( |
True
|
instance-attribute
¶
observation_space: Space[ObsType, BDeviceType, BDtypeType, BRNGType]
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
send_command(index: int, cmd: Literal['reset', 'step', 'render', 'close', '_describe', '_shared_setup'], *args, **kwargs)
Send a command to one worker process.
get_command_result(index: int)
Receive one worker result and surface worker-side exceptions.
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(*args, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> None
Dispatch reset commands to all selected workers.
reset_wait() -> Tuple[ContextType, ObsType, Dict[str, Any]]
Collect the pending reset results and batch them.
step(action: ActType) -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]
Convenience wrapper around step_async plus step_wait.
step_async(action: ActType) -> None
Dispatch one action batch to all worker processes.
With shared actions the whole batch is written into the shared buffers at once and the workers receive a payload-free command.
step_wait() -> Tuple[ObsType, BArrayType, BArrayType, BArrayType, Dict[str, Any]]
Collect the pending step results and batch them.
render() -> Sequence[RenderFrame] | None
Collect render outputs from every worker that returns a frame.
close()
Shut down all worker processes and communication pipes, then release the shared observation buffers (if any). Safe to call more than once.
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.
Space-forwarding invariant¶
By default the wrapper exposes the wrapped env's action_space,
observation_space, and context_space unchanged. Because the
:class:Env contract guarantees that these spaces already include the
batch dimension when the env is batched, subclasses that do override
a space must ensure the replacement still matches the batching status
of the underlying env. Never wrap an already-batched env space in an
additional batch_space call.
Wrap env and optionally override spaces or metadata later.
property
writable
¶
action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]
property
¶
unwrapped: Env
Returns the base environment of the wrapper.
This will be the bare :class:gymnasium.Env environment, underneath all layers of wrappers.
property
¶
prev_wrapper_layer: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
step(action: WrapperActT) -> Tuple[WrapperObsT, Union[SupportsFloat, WrapperBArrayT], Union[bool, WrapperBArrayT], Union[bool, WrapperBArrayT], Dict[str, Any]]
reset(*args, mask: Optional[WrapperBArrayT] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]
has_wrapper_attr(name: str) -> bool
Checks if the given attribute is within the wrapper or its environment.
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(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(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.
property
writable
¶
action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]
property
¶
unwrapped: Env
Returns the base environment of the wrapper.
This will be the bare :class:gymnasium.Env environment, underneath all layers of wrappers.
property
¶
prev_wrapper_layer: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]
reset(*args, mask: Optional[WrapperBArrayT] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the given attribute is within the wrapper or its environment.
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(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 |
abstractmethod
¶
map_action(action: WrapperActT) -> ActType
Map an outer action into the wrapped environment's action space.
reverse_map_action(action: ActType) -> WrapperActT
Map an inner action back to the wrapper-facing representation.
step(action: WrapperActT) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]
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.
property
writable
¶
action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]
property
¶
unwrapped: Env
Returns the base environment of the wrapper.
This will be the bare :class:gymnasium.Env environment, underneath all layers of wrappers.
property
¶
prev_wrapper_layer: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the given attribute is within the wrapper or its environment.
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(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(context: WrapperContextT) -> ContextType
Convert wrapped context into the wrapper-facing representation.
reverse_map_context(context: ContextType) -> WrapperContextT
Convert wrapper-facing context back into the wrapped representation.
map_observation(observation: WrapperObsT) -> ObsType
Convert wrapped observations into the wrapper-facing representation.
reverse_map_observation(observation: ObsType) -> WrapperObsT
Convert wrapper-facing observations back into the wrapped representation.
reset(*args, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperContextT, WrapperObsT, Dict[str, Any]]
step(action: ActType) -> Tuple[WrapperObsT, Union[float, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]
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.
Batched-environment space invariant¶
The same invariant as :class:Env applies: when batch_size is not
None (including batch_size == 1), the observation_space,
action_space, and context_space already include the leading
batch dimension. Callers, wrappers, and downstream code must not
apply an additional batch_space on top of these spaces.
instance-attribute
¶
observation_space: Space[Any, BDeviceType, BDtypeType, BRNGType]
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
abstractmethod
¶
initial(*, seed: Optional[int] = None, **kwargs) -> Tuple[StateType, ContextType, ObsType, Dict[str, Any]]
Create the first environment state, context, observation, and info.
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
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(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(state: StateType, render_state: RenderStateType) -> Tuple[RenderFrame | Sequence[RenderFrame] | None, StateType, RenderStateType]
Produce a render frame and updated runtime state.
render_close(state: StateType, render_state: RenderStateType) -> StateType
Release render-specific resources and return the updated state.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
property
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]]
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
step(action: ActType) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]
reset(*, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs) -> Tuple[ContextType, ObsType, Dict[str, Any]]
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.
Space-forwarding invariant¶
By default the wrapper forwards the wrapped FuncEnv's
action_space, observation_space, and context_space unchanged.
Because the :class:FuncEnv contract guarantees that these spaces
already include the batch dimension when the env is batched, subclasses
that override a space must ensure the replacement still matches the
batching status of the underlying env. Never wrap an already-batched
env space in an additional batch_space call.
Wrap a functional environment and optionally override its interface metadata.
property
writable
¶
action_space: Space[WrapperActT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
observation_space: Space[WrapperObsT, WrapperBDeviceT, Any, WrapperBRngT]
property
writable
¶
context_space: Optional[Space[WrapperContextT, WrapperBDeviceT, Any, WrapperBRngT]]
property
¶
prev_wrapper_layer: FuncEnv[StateType, RenderStateType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
initial(*, seed: Optional[int] = None, **kwargs) -> Tuple[WrapperStateT, WrapperContextT, WrapperObsT, Dict[str, Any]]
reset(state: WrapperStateT, *args, seed: Optional[int] = None, mask: Optional[WrapperBArrayT] = None, **kwargs) -> Tuple[WrapperStateT, WrapperContextT, WrapperObsT, Dict[str, Any]]
step(state: WrapperStateT, action: WrapperActT) -> Tuple[WrapperStateT, WrapperObsT, Union[SupportsFloat, WrapperBArrayT], Union[bool, WrapperBArrayT], Union[bool, WrapperBArrayT], Dict[str, Any]]
render_init(state: WrapperStateT, *, seed: Optional[int] = None, render_mode: Optional[str] = None, **kwargs) -> Tuple[WrapperStateT, WrapperRenderStateT, FuncEnvCommonRenderInfo]
render_image(state: WrapperStateT, render_state: WrapperRenderStateT) -> Tuple[WrapperRenderFrame | Sequence[WrapperRenderFrame] | None, WrapperStateT, WrapperRenderStateT]
render_close(state: WrapperStateT, render_state: WrapperRenderStateT) -> WrapperStateT
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
instance-attribute
¶
device: Optional[BDeviceType]
The world timestep in seconds, if None, the world is asynchronous (real-time)
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
class-attribute
instance-attribute
¶
world_subtimestep: Optional[float] = None
The number of parallel environments in this world
abstractmethod
¶
step() -> Union[float, BArrayType]
Step the world by one timestep. Returns: float: The elapsed time since the last step.
abstractmethod
¶
reset(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None
Reset the world state in-place.
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(*, 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(*, 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(control_timestep: Optional[float]) -> bool
Check whether a node control period aligns with world_timestep.
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(*, 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(*, 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(*, 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(control_timestep: Optional[float]) -> bool
Check whether a node control period aligns with world_timestep.
reset(*, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None
Start a new wall-clock timing sequence.
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)
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
supported_render_modes: Sequence[str] = ()
class-attribute
instance-attribute
¶
world: Optional[World[BArrayType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
after_reset_priorities: Set[int] = set()
class-attribute
instance-attribute
¶
after_reload_priorities: Set[int] = set()
class-attribute
instance-attribute
¶
pre_environment_step_priorities: Set[int] = set()
class-attribute
instance-attribute
¶
post_environment_step_priorities: Set[int] = set()
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
property
¶
effective_update_timestep: Optional[float]
Resolved update period, defaulting to control_timestep when unset.
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() -> Optional[ContextType]
Get the current context from the node. If the context space is None, this method should not be called.
get_observation() -> ObsType
Get the current observation from the sensor. If the observation space is None, this method should not be called.
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() -> 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() -> 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() -> Optional[Dict[str, Any]]
Get optional auxiliary information with this node.
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(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(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(*, 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(*, 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(*, 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(*, 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(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(fn: Callable[[WorldNode], bool]) -> list[WorldNode]
Return nodes in this subtree matching fn.
Vanilla nodes only evaluate self.
get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
Bases: ABC, Generic[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Functional counterpart to :class:World with explicit state passing.
instance-attribute
¶
device: Optional[BDeviceType]
The world timestep in seconds. If None, the world is asynchronous (real-time)
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
class-attribute
instance-attribute
¶
world_subtimestep: Optional[float] = None
The number of parallel environments in this world
abstractmethod
¶
initial(*, seed: Optional[int] = None, **kwargs) -> WorldStateT
Construct the initial world state.
abstractmethod
¶
step(state: WorldStateT) -> Tuple[WorldStateT, Union[float, BArrayType]]
Advance state by one world step.
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(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(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(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(control_timestep: Optional[float]) -> bool
Check whether a node control period aligns with world_timestep.
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.
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
supported_render_modes: Sequence[str] = ()
class-attribute
instance-attribute
¶
world: Optional[FuncWorld[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
after_reset_priorities: Set[int] = set()
class-attribute
instance-attribute
¶
after_reload_priorities: Set[int] = set()
class-attribute
instance-attribute
¶
pre_environment_step_priorities: Set[int] = set()
class-attribute
instance-attribute
¶
post_environment_step_priorities: Set[int] = set()
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
property
¶
effective_update_timestep: Optional[float]
Resolved update period, defaulting to control_timestep when unset.
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(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.
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(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(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(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(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(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(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(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(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(world_state: WorldStateT, node_state: NodeStateT) -> Optional[Dict[str, Any]]
Get the current info from the environment.
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(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(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(world_state: WorldStateT, node_state: NodeStateT) -> WorldStateT
Release node-owned resources and return the final world state.
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(fn: Callable[[FuncWorldNode], bool]) -> list[FuncWorldNode]
Return nodes in this subtree matching fn.
Vanilla nodes only evaluate self.
get_nodes_by_type(node_type: Type[FuncWorldNode]) -> list[FuncWorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
property
¶
control_timestep: Optional[float]
Simulation time represented by one environment control step.
This is the resolved control timestep: the explicit node
control_timestep when configured, otherwise the effective
update timestep, otherwise the world timestep. It is None
when none of these are defined (e.g. an asynchronous real-time
world).
sample_context() -> Optional[ContextType]
Sample one context value if context_space is defined.
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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(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(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(*, 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(*, 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(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() -> RenderFrame | Sequence[RenderFrame] | None
Render via the root node when rendering is enabled.
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.
property
¶
control_timestep: Optional[float]
Simulation time represented by one environment control step.
This is the resolved control timestep: the explicit node
control_timestep when configured, otherwise the effective
update timestep, otherwise the world timestep. It is None
when none of these are defined (e.g. an asynchronous real-time
world).
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(old_context: ContextType, new_context: ContextType, mask: BArrayType) -> ContextType
Merge masked reset contexts back into a full batched context.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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(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(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(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(*, seed: Optional[int] = None, **kwargs) -> Tuple[WorldFuncEnvState, ContextType, ObsType, Dict[str, Any]]
Create the initial world state and root node state.
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(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(state: WorldFuncEnvState) -> None
Close the node tree first, then the underlying world.
render_init(state, *, seed=None, render_mode=None, **kwargs)
Return a no-op render state for node-based functional environments.
render_image(state, render_state)
Render through the root node when rendering is enabled.
render_close(state, render_state)
Return the unchanged state because no extra render state is held.
dataclass
¶
WorldFuncEnvState(world_state: Any, node_state: Any)
EnvAsWorldNode(env: Env, name: str, *, world: Optional[World] = None, control_timestep: Optional[float] = None, update_timestep: Optional[float] = None, hold_last_action: bool = False, expose_reward: bool = True, expose_termination: bool = True, expose_truncation: bool = True, forward_info: bool = True, reset_seed: Optional[int] = None, pre_step_priority: int = 0)
Bases: WorldNode[ContextType, ObsType, ActType, BArrayType, BDeviceType, BDtypeType, BRNGType]
Adapt an existing unbatched :class:Env so it can be composed as a WorldNode.
Design¶
Env.step(action) is atomic: it produces the next observation, reward,
termination/truncation signals, and info in a single call. A WorldNode
instead splits the step lifecycle into::
set_next_action(action)
-> pre_environment_step(dt, priority=...)
-> World.step()
-> post_environment_step(dt, priority=...)
-> get_observation() / get_reward() / ... / get_info()
This adapter bridges the mismatch by performing the atomic inner
Env.step inside pre_environment_step (using the action cached by the
preceding set_next_action call) and caching its five-tuple result for
the get_* accessors to expose. reset maps directly to
Env.reset, and the base WorldNode.reload default (which delegates to
reset) is exactly right for a pre-built environment, hence the
non-empty reset_priorities / reload_priorities sets.
The dt argument of pre_environment_step is intentionally ignored:
the wrapped environment is self-paced (e.g. real hardware running at its own
control rate), while the outer world is expected to be a bookkeeping-only
world such as :class:RealWorld. control_timestep / update_timestep
can still be set for step-ratio bookkeeping by the composer.
No schema transformation / re-keying is performed: the inner environment's action, observation, and context spaces (and the raw observation / context layouts) pass through unchanged.
Only unbatched environments (batch_size is None) are supported. The
optional world is only used for backend/device delegation; when provided
its backend type must match the inner environment's backend, since the node
serves the inner environment's arrays.
property
¶
effective_update_timestep: Optional[float]
Resolved update period, defaulting to control_timestep when unset.
class-attribute
instance-attribute
¶
pre_environment_step_priorities = {pre_step_priority}
class-attribute
instance-attribute
¶
post_environment_step_priorities = set()
instance-attribute
¶
update_timestep = update_timestep if update_timestep is not None else control_timestep
instance-attribute
¶
supported_render_modes = tuple(env.metadata.get('render_modes', ()))
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend of the attached world, falling back to the inner env's backend.
property
¶
device: Optional[BDeviceType]
Device of the attached world, falling back to the inner env's device.
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.
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(*, 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(*, 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(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(fn: Callable[[WorldNode], bool]) -> list[WorldNode]
Return nodes in this subtree matching fn.
Vanilla nodes only evaluate self.
get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
set_next_action(action: ActType) -> None
Cache the action consumed by the next pre_environment_step call.
action=None is only accepted when hold_last_action=True (the
previously cached action is kept); in that case a node that has never
received an action still raises from pre_environment_step.
pre_environment_step(dt: Union[float, BArrayType], *, priority: int = 0) -> None
Run one atomic Env.step with the cached action.
The world-provided dt is ignored: the wrapped environment is
self-paced, so it advances by exactly one of its own steps per call.
reset(*, priority: int = 0, seed: Optional[int] = None, mask: Optional[BArrayType] = None, **kwargs) -> None
Reset the inner environment (also used as the default reload).
get_context() -> Optional[ContextType]
Return the context produced by the most recent inner reset.
get_observation() -> ObsType
Return the latest step observation, or the reset observation before the first step.
get_reward() -> Union[float, BArrayType]
Return the cached step reward (0.0 before the first step).
get_termination() -> Union[bool, BArrayType]
Return the cached termination signal (False before the first step).
get_truncation() -> Union[bool, BArrayType]
Return the cached truncation signal (False before the first step).
get_info() -> Optional[Dict[str, Any]]
Return a copy of the latest info dict when forward_info is enabled.
CombinedWorldNode(name: str, nodes: Iterable[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], direct_return: bool = True, render_mode: Optional[str] = 'auto', root_node: Optional[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None)
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.
Root node¶
An optional root_node may be passed next to nodes. It is a regular
node (sharing the same :class:~unienv_interface.world.world.World) that is
NOT part of the public nodes list and whose data merges at the TOP LEVEL
of the aggregated dictionaries instead of nesting under its own name. It
rides the same machinery as a child node -- priorities, control/update rate
ratios, lifecycle dispatch, rewards/signals, info, render and close -- under
the reserved internal key :data:ROOT_NODE_KEY (""), which cannot
collide with a child node name (child names are validated non-empty).
- Spaces / data (context, observation, action -- independently). When the
root node is the sole contributor of a channel, the existing
direct_returnrule applies verbatim:direct_return=Truepasses the root's own space/data through unwrapped (anySpacetype, e.g. a plainBoxSpace), whiledirect_return=Falsewraps it under the reserved key (DictSpace({ROOT_NODE_KEY: space})/{ROOT_NODE_KEY: data}). When the root and at least one named child contribute, the root MUST expose aDictSpacefor that channel and its keys are merged next to the child entries:{child_name: child_space, ..., **root_space.spaces}(resp.{child_name: child_data, ..., **root_data}). A root key equal to a child node name raisesValueError. - Actions. For a sole-root action channel the aggregate action space is the root's action space and the whole action value is forwarded to the root at its own control-rate tick. In the mixed case the incoming mapping is routed by key: keys matching action-child names update those children's caches (exactly as before), and every remaining key accumulates into the root's action entry as a dict. Each consumer is dispatched at its own ratio tick with hold-last semantics identical to children.
- Info / render. These channels nest by node name for children, so the
root's entries are keyed by the root's own
name(the reserved key is only used for internal state/routing and aggregation). get_nodenever resolves the reserved key (it is not a user-facing path);get_nodes_by_fndoes traverse the root node.
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
class-attribute
instance-attribute
¶
supported_render_modes = ('dict', 'auto')
instance-attribute
¶
has_termination_signal = any(node.has_termination_signal for node in self._all_nodes)
instance-attribute
¶
has_truncation_signal = any(node.has_truncation_signal for node in self._all_nodes)
get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
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_priority_order(phase: str) -> List[int]
Return the cached descending priority order for a named lifecycle phase.
Supported phase names: 'reload', 'after_reload', 'reset',
'after_reset', 'pre_step', 'post_step'.
get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[WorldNode]
Resolve a child node by dotted-path-like key traversal.
get_nodes_by_fn(fn: Callable[[WorldNode], bool]) -> list[WorldNode]
Return every node in the subtree that satisfies fn (root included).
pre_environment_step(dt, *, priority: int = 0)
Dispatch pre-step callbacks to the nodes (children and the optional root node) at the matching frequency.
set_next_action(action)
Route combined actions to child nodes (and the optional root node), respecting per-node control rates.
The incoming action is split by :meth:_split_child_actions (children
only). When a root_node owns an action space and no single node is
the direct-action passthrough provider, the keys that do not belong to
any named child are the root's action keys: they accumulate into the
root's cache entry under the reserved key as a dict. Every consumer
(children and root) is then dispatched at its own ratio tick, and only
once it has received at least one action (hold-last semantics, so a
child-only step re-delivers the root's previous action at its tick and
vice versa).
post_environment_step(dt, *, priority: int = 0)
Dispatch post-step callbacks to the nodes (children and the optional root node) at the matching frequency.
reset(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})
Forward reset calls to the nodes (children and the optional root node)
that participate at priority.
reload(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})
Forward reload calls to the nodes (children and the optional root node)
that participate at priority.
after_reset(*, priority: int = 0, mask=None)
Forward post-reset hooks and reset internal routing counters.
after_reload(*, priority: int = 0, mask=None)
Call after_reload on the nodes (children and the optional root node). Similar to after_reset but for the reload flow.
FlatCombinedWorldNode(name: str, nodes: Iterable[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], render_mode: Optional[str] = 'auto', root_node: Optional[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None)
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.
Root node¶
An optional root_node (passed separately, sharing the same world and not
listed in nodes) is handled as one more flat contributor:
- Its context / observation spaces and data are appended to the flatten list,
so its keys appear directly at the top level next to the children's keys
(an overlapping key raises
ValueError). When the root is the only contributor its space/data is returned unwrapped, mirroring the existing single-child flat behavior. - For actions, every key not claimed by any child's action space is routed to
the root under the reserved internal key, while claimed keys are sliced to
their children exactly as before. Without a
root_nodesuch unclaimed keys are rejected with aValueError. - Rewards / signals / info / render and the lifecycle hooks include the root
like any other node (info and render use the root's own
name).
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'
|
root_node
|
Optional[WorldNode[Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]
|
Optional node merged at the top level of the flattened data instead of nesting under its own name |
None
|
Raises:
| Type | Description |
|---|---|
ValueError
|
If node spaces have overlapping keys or non-DictSpace types |
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None
instance-attribute
¶
has_termination_signal = any(node.has_termination_signal for node in self._all_nodes)
instance-attribute
¶
has_truncation_signal = any(node.has_truncation_signal for node in self._all_nodes)
class-attribute
instance-attribute
¶
supported_render_modes = ('dict', 'auto')
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
pre_environment_step(dt, *, priority: int = 0)
Dispatch pre-step callbacks to the nodes (children and the optional root node) at the matching frequency.
set_next_action(action)
Route combined actions to child nodes (and the optional root node), respecting per-node control rates.
The incoming action is split by :meth:_split_child_actions (children
only). When a root_node owns an action space and no single node is
the direct-action passthrough provider, the keys that do not belong to
any named child are the root's action keys: they accumulate into the
root's cache entry under the reserved key as a dict. Every consumer
(children and root) is then dispatched at its own ratio tick, and only
once it has received at least one action (hold-last semantics, so a
child-only step re-delivers the root's previous action at its tick and
vice versa).
post_environment_step(dt, *, priority: int = 0)
Dispatch post-step callbacks to the nodes (children and the optional root node) at the matching frequency.
reset(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})
Forward reset calls to the nodes (children and the optional root node)
that participate at priority.
reload(*, priority: int = 0, seed=None, mask=None, pernode_kwargs: Dict[str, Any] = {})
Forward reload calls to the nodes (children and the optional root node)
that participate at priority.
after_reset(*, priority: int = 0, mask=None)
Forward post-reset hooks and reset internal routing counters.
after_reload(*, priority: int = 0, mask=None)
Call after_reload on the nodes (children and the optional root node). Similar to after_reset but for the reload flow.
get_node(nested_keys: Union[str, Sequence[str]]) -> Optional[WorldNode]
Resolve a child node by dotted-path-like key traversal.
get_nodes_by_fn(fn: Callable[[WorldNode], bool]) -> list[WorldNode]
Return every node in the subtree that satisfies fn (root included).
get_nodes_by_type(node_type: Type[WorldNode]) -> list[WorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
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_priority_order(phase: str) -> List[int]
Return the cached descending priority order for a named lifecycle phase.
Supported phase names: 'reload', 'after_reload', 'reset',
'after_reset', 'pre_step', 'post_step'.
get_context() -> Optional[CombinedDataT]
Get context by flattening all node contexts (root node included) into one dictionary.
get_observation() -> CombinedDataT
Get observation by flattening all node observations (root node included) into one dictionary.
CombinedFuncWorldNode(name: str, nodes: Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], direct_return: bool = True, render_mode: Optional[str] = 'auto', root_node: Optional[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None)
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.
Root node¶
An optional root_node may be passed next to nodes. It is a regular
FuncWorldNode (sharing the same FuncWorld) that is NOT part of the
public nodes list and whose data merges at the TOP LEVEL of the aggregated
dictionaries instead of nesting under its own name. It rides the same
machinery as a child node -- priorities, control/update rate ratios, lifecycle
dispatch, rewards/signals, info, render and close -- and its node state is
stored in the combined node_state under the reserved key
:data:ROOT_NODE_KEY (""), a string key that cannot collide with a child
node name (child names are validated non-empty) and therefore keeps persisted
CombinedNodeStateT dictionaries loadable without migration.
- Spaces / data (context, observation, action -- independently). Spaces are
aggregated once at construction (functional nodes declare static spaces).
When the root node is the sole contributor of a channel, the existing
direct_returnrule applies verbatim:direct_return=Truepasses the root's own space/data through unwrapped (anySpacetype, e.g. a plainBoxSpace), whiledirect_return=Falsewraps it under the reserved key (DictSpace({ROOT_NODE_KEY: space})/{ROOT_NODE_KEY: data}). When the root and at least one named child contribute, the root MUST expose aDictSpacefor that channel and its keys are merged next to the child entries:{child_name: child_space, ..., **root_space.spaces}(resp.{child_name: child_data, ..., **root_data}). A root key equal to a child node name raisesValueError. - Actions. For a sole-root action channel the aggregate action space is the root's action space and the whole action value is forwarded to the root at its own control-rate tick. In the mixed case the incoming mapping is routed by key: keys matching action-child names update those children's cache entries (exactly as before) and every remaining key accumulates into the root's entry (stored under the reserved key) as a dict. Each consumer is dispatched at its own ratio tick with hold-last semantics identical to children.
- Info / render. These channels nest by node name for children, so the
root's entries are keyed by the root's own
name(the reserved key is only used for internal state/routing and aggregation). get_nodenever resolves the reserved key (it is not a user-facing path);get_nodes_by_fndoes traverse the root node.
class-attribute
instance-attribute
¶
context_space: Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
observation_space: Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
class-attribute
instance-attribute
¶
action_space: Optional[Space[ActType, BDeviceType, BDtypeType, BRNGType]] = None
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
class-attribute
instance-attribute
¶
supported_render_modes = ('dict', 'auto')
instance-attribute
¶
has_termination_signal = any(node.has_termination_signal for node in self._all_nodes)
instance-attribute
¶
has_truncation_signal = any(node.has_truncation_signal for node in self._all_nodes)
get_nodes_by_type(node_type: Type[FuncWorldNode]) -> list[FuncWorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
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(nested_keys: Union[str, Sequence[str]]) -> Optional[FuncWorldNode]
Resolve a child node by dotted-path-like key traversal.
get_nodes_by_fn(fn: Callable[[FuncWorldNode], bool]) -> list[FuncWorldNode]
Return every node in the subtree that satisfies fn (root included).
initial(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]
Create node states (children and the optional root node, whose state is stored under the reserved key) for the current initialization priority.
reload(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]
Recreate node states (children and the optional root node) for the current reload priority.
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 the nodes (children and the optional root node)
that participate at priority.
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(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(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]
set_next_action(world_state: WorldStateT, node_state: CombinedNodeStateT, action: CombinedDataT) -> Tuple[WorldStateT, CombinedNodeStateT]
Route combined actions to child nodes (and the optional root node), respecting per-node control rates.
The incoming action is split by :meth:_split_child_actions (children
only). When a root_node owns an action space and no single node is the
direct-action passthrough provider, the keys that do not belong to any
named child are the root's action keys: they accumulate into the root's
cache entry (stored under the reserved key) as a dict. Every consumer
(children and root) is then dispatched at its own ratio tick, and only once
it has received at least one action (hold-last semantics, so a child-only
step re-delivers the root's previous action at its tick and vice versa).
post_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]
close(world_state: WorldStateT, node_state: CombinedNodeStateT) -> WorldStateT
Close every child node and the optional root node.
get_context(world_state: WorldStateT, node_state: CombinedNodeStateT) -> CombinedDataT
get_observation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> CombinedDataT
get_reward(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[float, BArrayType]
get_termination(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]
get_truncation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]
get_info(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Optional[Dict[str, Any]]
FlatCombinedFuncWorldNode(name: str, nodes: Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], render_mode: Optional[str] = 'auto', root_node: Optional[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None)
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.
Root node¶
An optional root_node (passed separately, sharing the same world and not
listed in nodes) is handled as one more flat contributor:
- Its context / observation spaces and data are appended to the flatten list,
so its keys appear directly at the top level next to the children's keys (an
overlapping key raises
ValueError). When the root is the only contributor its space/data is returned unwrapped, mirroring the existing single-child flat behavior. - For actions, every key not claimed by any child's action space is routed to
the root under the reserved internal key (its state lives under that same
key), while claimed keys are sliced to their children exactly as before.
Without a
root_nodesuch unclaimed keys are rejected with aValueError. - Rewards / signals / info / render and the lifecycle hooks include the root
like any other node (info and render use the root's own
name).
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'
|
root_node
|
Optional[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]]
|
Optional node merged at the top level of the flattened data instead of nesting under its own name |
None
|
Raises:
| Type | Description |
|---|---|
ValueError
|
If node spaces have overlapping keys or non-DictSpace types |
instance-attribute
¶
has_termination_signal = any(node.has_termination_signal for node in self._all_nodes)
instance-attribute
¶
has_truncation_signal = any(node.has_truncation_signal for node in self._all_nodes)
class-attribute
instance-attribute
¶
supported_render_modes = ('dict', 'auto')
property
¶
backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
Backend provided by the attached world.
instance-attribute
¶
observation_space = self._flatten_channel_spaces('observation_space')
initial(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]
Create node states (children and the optional root node, whose state is stored under the reserved key) for the current initialization priority.
reload(world_state: WorldStateT, *, priority: int = 0, seed: Optional[int] = None, pernode_kwargs: Dict[str, Dict[str, Any]] = {}) -> Tuple[WorldStateT, CombinedNodeStateT]
Recreate node states (children and the optional root node) for the current reload priority.
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 the nodes (children and the optional root node)
that participate at priority.
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(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(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]
get_reward(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[float, BArrayType]
get_termination(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]
get_truncation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Union[bool, BArrayType]
set_next_action(world_state: WorldStateT, node_state: CombinedNodeStateT, action: CombinedDataT) -> Tuple[WorldStateT, CombinedNodeStateT]
Route combined actions to child nodes (and the optional root node), respecting per-node control rates.
The incoming action is split by :meth:_split_child_actions (children
only). When a root_node owns an action space and no single node is the
direct-action passthrough provider, the keys that do not belong to any
named child are the root's action keys: they accumulate into the root's
cache entry (stored under the reserved key) as a dict. Every consumer
(children and root) is then dispatched at its own ratio tick, and only once
it has received at least one action (hold-last semantics, so a child-only
step re-delivers the root's previous action at its tick and vice versa).
post_environment_step(world_state: WorldStateT, node_state: CombinedNodeStateT, dt: Union[float, BArrayType], *, priority: int = 0) -> Tuple[WorldStateT, CombinedNodeStateT]
close(world_state: WorldStateT, node_state: CombinedNodeStateT) -> WorldStateT
Close every child node and the optional root 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(fn: Callable[[FuncWorldNode], bool]) -> list[FuncWorldNode]
Return every node in the subtree that satisfies fn (root included).
get_nodes_by_type(node_type: Type[FuncWorldNode]) -> list[FuncWorldNode]
Return nodes in this subtree that are instances of node_type.
has_wrapper_attr(name: str) -> bool
Checks if the attribute name exists in the environment.
set_wrapper_attr(name: str, value: Any)
Sets the attribute name on the environment with value.
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.
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(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Optional[CombinedDataT]
Get context by flattening all node contexts into one dictionary.
get_observation(world_state: WorldStateT, node_state: CombinedNodeStateT) -> CombinedDataT
Get observation by flattening all node observations into one dictionary.
get_info(world_state: WorldStateT, node_state: CombinedNodeStateT) -> Optional[Dict[str, Any]]
Get info by merging all node info dictionaries.
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.
Reassigning an attribute of a transformation increments the global
DataTransformation._config_epoch (attributes whose name ends with
_cache are exempt), so caches derived from a transformation's
configuration (e.g. ChainedTransformation's intermediate spaces) are
rebuilt. Containers and arrays held by a transformation should be replaced,
not mutated in place, once it has been used.
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.
abstractmethod
¶
transform(source_space: Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT], data: SourceDataT) -> TargetDataT
Apply the transformation to data drawn from source_space.
direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]
Return the inverse transformation when one exists.
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 |
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 |
Bases: DataTransformation
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(enabled_keys: Iterable[str], *, ignore_missing_keys: bool = False, nested_separator: str = '/')
Bases: DataTransformation
direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]
Return the inverse transformation when one exists.
serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> DictIncludeKeyTransformation
DictExcludeKeyTransformation(excluded_keys: Iterable[str], *, ignore_missing_keys: bool = False, nested_separator: str = '/')
Bases: DataTransformation
direction_inverse(source_space: Optional[Space[SourceDataT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT]] = None) -> Optional[DataTransformation]
Return the inverse transformation when one exists.
serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> DictExcludeKeyTransformation
BatchifyTransformation(axis: int = 0)
Bases: DataTransformation
UnBatchifyTransformation(axis: int = 0)
Bases: DataTransformation
DictTransformation(mapping: Dict[str, DataTransformation], ignore_missing_keys: bool = False, nested_separator: str = '/')
Bases: DataTransformation
instance-attribute
¶
has_inverse = all(transformation.has_inverse for transformation in mapping.values())
get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType])
transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> DictTransformation
FlattenDictTransformation(nested_separator: str = '/')
Bases: DataTransformation
get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType]) -> DictSpace[BDeviceType, BDtypeType, BRNGType]
transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]
direction_inverse(source_space: Optional[Space] = None) -> UnflattenDictTransformation
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> FlattenDictTransformation
UnflattenDictTransformation(nested_separator: str = '/')
Bases: DataTransformation
get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType]) -> DictSpace[BDeviceType, BDtypeType, BRNGType]
transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]
direction_inverse(source_space: Optional[Space] = None) -> FlattenDictTransformation
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> UnflattenDictTransformation
ChainedTransformation(transformations: Iterable[DataTransformation])
Bases: DataTransformation
instance-attribute
¶
has_inverse = all(transformation.has_inverse for transformation in self.transformations)
get_target_space_from_source(source_space: DictSpace[BDeviceType, BDtypeType, BRNGType])
direction_inverse(source_space: Optional[Space] = None) -> Optional[ChainedTransformation]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space] = None) -> ChainedTransformation
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(source_space: Space[Any, BDeviceType, BDtypeType, BRNGType]) -> None
get_crop_range(source_space: BoxSpace[Any, BDeviceType, BDtypeType, BRNGType]) -> Tuple[BArrayType, BArrayType]
serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> CropTransformation
ImageResizeTransformation(new_height: int, new_width: int)
Bases: DataTransformation
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(source_space: Space[Any, BDeviceType, BDtypeType, BRNGType])
transform(source_space: Space, data: Union[Mapping[str, Any], BArrayType]) -> Union[Mapping[str, Any], BArrayType]
direction_inverse(source_space=None) -> Optional[Union[IterativeTransformation, LeafwiseTransformation]]
Return the inverse transformation.
Inverses that depend on the source space (e.g. RescaleTransformation)
are computed per leaf of source_space. When all leaves are equal a
single IterativeTransformation is returned, otherwise a
LeafwiseTransformation holding one inverse per leaf.
serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> IterativeTransformation
LeafwiseTransformation(transformations: Union[DataTransformation, Mapping[str, Any], Tuple[Any, ...]])
Bases: DataTransformation
Apply a separate transformation to each leaf of a dict/tuple space.
transformations mirrors the space structure: a mapping for a
DictSpace, a tuple for a TupleSpace and a DataTransformation for
a leaf. This is what :meth:IterativeTransformation.direction_inverse
returns when the leaves need different inverses (e.g. rescaling leaves with
different source bounds).
instance-attribute
¶
has_inverse = all(t.has_inverse for t in _iter_leaves(transformations))
serialize(source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> Dict[str, Any]
classmethod
¶
deserialize_from(json_data: Dict[str, Any], source_space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]] = None) -> LeafwiseTransformation
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_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. |