unienv_interface.env_base.vec_env¶
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)
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.
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.
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
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.
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.
data_stack(data: Any, backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], device: Optional[BDeviceType] = None)
Recursively stack per-environment data into a batched backend representation.
- A mapping is processed value by value.
- A sequence of mappings that share the same keys (e.g. one
infodict per sub-environment) becomes a mapping of stacked values. - A sequence of arrays or Python numbers is stacked into one array.
- Strings, bytes and other objects are kept as a list.