Skip to content

unienv_interface.space.space_utils

reshape_batch_size

reshape_batch_size(space: Space, old_batch_shape: Tuple[int], new_batch_shape: Tuple[int]) -> Space

reshape_batch_size_in_data

reshape_batch_size_in_data(backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], data: Any, old_batch_shape: Tuple[int], new_batch_shape: Tuple[int]) -> Any

swap_batch_dims

swap_batch_dims(space: Space, dim1: int, dim2: int) -> Space

swap_batch_dims_in_data

swap_batch_dims_in_data(backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], data: Any, dim1: int, dim2: int) -> Any

batch_size

batch_size(space: Space) -> Optional[int]

Space does not support batching.

batch_space

batch_space(space: Space, n: int = 1) -> Space

batch_differing_spaces

batch_differing_spaces(spaces: Sequence[Space], device: Optional[Any] = None) -> Space

unbatch_spaces

unbatch_spaces(space: Space) -> Iterable[Space]

iterate

iterate(space: Space, items: Any) -> Iterator

get_at

get_at(space: Space, items: Any, index: ArrayAPIGetIndex) -> Any

set_at

set_at(space: Space, items: Any, index: ArrayAPISetIndex, value: Any) -> Any

concatenate

concatenate(space: Space, items: Iterable[Any], axis: int = 0) -> Any

construct_space_from_data_stream

construct_space_from_data_stream(data: Iterable[Any], backend: ComputeBackend, add_bounds: bool = True) -> Space

Construct a space from a stream of data.

Parameters:

Name Type Description Default
data Iterable[Any]

An iterable stream of data samples.

required
backend ComputeBackend

The compute backend to use for array operations.

required

Returns:

Name Type Description
Space Space

The constructed space.

construct_space_from_data

construct_space_from_data(data: Any, backend: ComputeBackend, candidate_space: Optional[Space] = None, add_bounds: bool = False) -> Space

is_flattenable

is_flattenable(space: Space, start_dim: int = 0) -> bool

flat_dim

flat_dim(space: Space, start_dim: int = 0) -> Optional[int]

flatten_space

flatten_space(space: Space, start_dim: int = 0) -> BoxSpace

Flatten space (from start_dim on) into a single BoxSpace.

The flat box carries no metadata: a description such as {"kind": "image", "layout": "HWC"} does not hold for the flattened (and, for structured spaces, concatenated) data.

flatten_data

flatten_data(space: Space, data: Any, start_dim: int = 0) -> BArrayType

unflatten_data

unflatten_data(space: Space, data: BArrayType, start_dim: int = 0) -> Any

space_to_json

space_to_json(space: Space) -> typing.Dict[str, Any]

Serialize a Space to a JSON-compatible dictionary.

The space's metadata (and that of every nested space) is stored under a "metadata" key, which is omitted when the metadata is empty.

Parameters:

Name Type Description Default
space Space

The space to serialize.

required

Returns:

Name Type Description
dict Dict[str, Any]

A JSON-compatible representation of the space.

json_to_space

json_to_space(json_data: Dict[str, Any], map_backend: ComputeBackend, map_device: Optional[BDeviceType]) -> Space

Deserialize a JSON-compatible dictionary into a Space instance.

Parameters:

Name Type Description Default
json_data dict

JSON-compatible space payload produced by :func:space_to_json.

required
map_backend ComputeBackend

Backend used to reconstruct backend-owned arrays and dtypes in the returned space.

required
map_device Optional[BDeviceType]

Optional device on which to place backend arrays created during deserialization.

required

Returns:

Name Type Description
Space Space

The deserialized space.

child_spaces

child_spaces(space: Space) -> List[Tuple[Union[str, int], Space]]

The nested spaces of space as (key, subspace) pairs.

Keys are dict keys for DictSpace, positions for TupleSpace / UnionSpace, and "single_space" / "node_feature_space" / "edge_feature_space" for BatchedSpace / GraphSpace.

spaces_equal_with_metadata

spaces_equal_with_metadata(a: Space, b: Space) -> bool

a == b and every pair of corresponding (nested) spaces has equal metadata.

Space.__eq__ ignores metadata by design; use this where the description of the data matters as well.

common_metadata

common_metadata(spaces: Iterable[Space]) -> SpaceMetadata

The metadata shared by all spaces: theirs if all are equal, else (or without spaces) empty.

Used where several spaces become one, e.g. batching differing spaces or merging the keys of several DictSpace s into one: a description that not every contributor agrees on fits none of the result.

find_spaces_by_metadata

find_spaces_by_metadata(space: Space, predicate: Optional[Callable[[Space], bool]] = None, **match: Any) -> List[Tuple[SpacePath, Space]]

Find the leaves of a (nested) DictSpace / TupleSpace by their metadata.

Returns (path, leaf) pairs, in iteration order, for every leaf (a space that is not a DictSpace / TupleSpace) whose metadata contains all match items and, if given, for which predicate(leaf) holds. path is the tuple of dict keys / tuple positions leading to the leaf (() when space itself is a leaf), i.e. the same path that indexes the data.

Example::

cameras = find_spaces_by_metadata(obs_space, kind="image")
# [(("cameras", "front"), BoxSpace(...)), (("wrist",), BoxSpace(...))]

metadata

metadata(kind: Optional[str], **fields: Any) -> Dict[str, Any]

A metadata dict of the given kind with fields; None values are dropped.

Example::

metadata("force", units="N", frame="wrist")
# {"kind": "force", "units": "N", "frame": "wrist"}

image_metadata

image_metadata(*, format: Optional[str] = 'rgb', layout: Optional[str] = 'HWC', source: Optional[str] = None, fps: Optional[float] = None, **extra: Any) -> Dict[str, Any]

Metadata of a color image (kind="image"); None values are dropped.

format is "rgb", "bgr", "rgba", "bgra" or "gray" and layout the order of the trailing axes, "HWC", "CHW" or "HW".

depth_metadata

depth_metadata(*, units: Optional[str] = 'm', layout: Optional[str] = 'HW', source: Optional[str] = None, fps: Optional[float] = None, **extra: Any) -> Dict[str, Any]

Metadata of a depth image (kind="depth"); None values are dropped.

layout is "HW" (no channel axis), or "HWC" / "CHW" with a single channel.

proprio_metadata

proprio_metadata(*, units: Optional[Union[str, Sequence[str]]] = None, names: Optional[Sequence[str]] = None, frame: Optional[str] = None, rotation: Optional[str] = None, **extra: Any) -> Dict[str, Any]

Metadata of proprioceptive data (kind="proprio"); None values are dropped.

units and names may be lists with one entry per element of the last axis.

action_metadata

action_metadata(*, control: Optional[str] = None, mode: Optional[str] = None, units: Optional[Union[str, Sequence[str]]] = None, normalized: Optional[bool] = None, range: Optional[Sequence[Any]] = None, names: Optional[Sequence[str]] = None, frame: Optional[str] = None, rotation: Optional[str] = None, **extra: Any) -> Dict[str, Any]

Metadata of an action (kind="action"); None values are dropped.

control is what is commanded (e.g. "joint_position", "eef_pose"), mode how ("absolute", "delta" or "velocity"). With normalized=True, range=[low, high] gives the physical range the normalized values map to.

annotate

annotate(space: Space, mapping: Mapping[Union[str, Tuple[Union[str, int], ...]], Mapping[str, Any]], *, replace: bool = False, strict: bool = True) -> Space

Return space with metadata attached to the nested spaces named in mapping.

Each key of mapping is a path through nested DictSpace and TupleSpace spaces: a "/"-separated string ("cameras/front", "0/joints"; tuple positions as numbers) or a tuple of dict keys / tuple positions (("cameras", "front")); "" or () is space itself. Its value is merged into the metadata of that space, or replaces it with replace=True (see :meth:Space.with_metadata).

Spaces that are not on an annotated path are kept as they are (the same objects), as is the metadata of every space that is not annotated; the containers on the paths are rebuilt. With strict=True (the default) a KeyError listing every path that does not exist is raised, otherwise such paths are skipped.

Example::

obs_space = from_gym_space(env.observation_space, backend)
obs_space = annotate(obs_space, {
    "pixels": image_metadata(source="front_cam"),
    "agent/qpos": proprio_metadata(units="rad"),
})

metadata_issues

metadata_issues(space: Space) -> List[str]

Problems with the metadata of space and all its nested spaces, as messages.

Checks the keys of the shared vocabulary (see "Metadata" in the Space docstring), leniently: unknown keys and kinds are fine and only what can be checked is checked.

  • types: kind, frame, control, mode, source and label are strings, units a string or a list of strings, names a list of strings, normalized a bool, fps a positive number and range [low, high] of numbers or lists of numbers (low <= high);
  • values: format ("rgb", "bgr", "rgba", "bgra", "gray"), layout ("HWC", "CHW", "HW") and rotation ("quat_wxyz", "quat_xyzw", "rot6d", "matrix", "axis_angle" or "euler_<axes>" with transforms3d axes such as "sxyz" / "rzyx");
  • array spaces (BoxSpace, BinarySpace, DynamicBoxSpace): lists in units, names and range have one entry per element of the last axis; layout matches the trailing axes ("HWC" / "CHW": at least 3, "HW": at least 2; leading batch axes are allowed), the channel count matches format (rgb / bgr: 3, rgba / bgra: 4, gray: 1) and kind="depth" has a single channel; kind="image" / "depth" spaces have at least 2 axes.

Each message starts with the "/"-separated path of the space ("<root>" for space itself), e.g. "cameras/front: format 'rgb' needs 3 channel(s), ...".

validate_metadata

validate_metadata(space: Space) -> None

Raise ValueError listing every :func:metadata_issues of space (if there is any).

Meant for tests and construction-time checks, e.g. in an adaptor's test suite::

def test_observation_space_metadata():
    validate_metadata(make_env().observation_space)