Spaces

Utilities operating over complex Gym.Space`s associated with arbitrarily nested data structure of `np.ndarray and heavily optimized for speed.

They combine static control flow pre-computed for given space and eventually some pre-allocated values with Just-In-Time (JIT) compiling via Numba when possible for optimal performance.

gym_jiminy.common.utils.spaces.get_bounds(space)[source]

Get the lower and upper bounds of a given ‘gym.Space’ if any.

Parameters:

space (Space) – gym.Space on which to operate.

Returns:

Lower and upper bounds as a tuple.

Return type:

Tuple[ndarray | SupportsFloat | None, ndarray | SupportsFloat | None]

gym_jiminy.common.utils.spaces.zeros(space, dtype=None, enforce_bounds=True)[source]

Allocate data structure from gym.Space and initialize it to zero.

Parameters:
Return type:

DataNestedT

gym_jiminy.common.utils.spaces.fill(data, fill_value)[source]

Set every element of ‘data’ from gym.Space to scalar ‘fill_value’.

Parameters:

  • data (Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray) – Data structure to update.

  • fill_value (float | int | number) – Value used to fill any scalar from the leaves.

Return type:

None

gym_jiminy.common.utils.spaces.copyto(dst, src)[source]

Copy arbitrarily nested data structure of ‘np.ndarray’ to a given pre-allocated destination.

It avoids memory allocation completely, so that memory pointers of ‘data’ remains unchanged. As direct consequences, it is necessary to preallocate memory beforehand, and it only supports arrays of fixed shape.

Note

Unlike the function returned by ‘build_copyto’, only the flattened data structure needs to match, not the original one. This means that the source and/or destination can be flattened already when provided. Beware values must be sorted by keys in case of nested dict.

Parameters:
Return type:

None

gym_jiminy.common.utils.spaces.copy(data)[source]

Shallow copy recursively ‘data’ from gym.Space, so that only leaves are still references.

Parameters:

data (DataNestedT) – Hierarchical data structure to copy without allocation.

Return type:

DataNestedT

gym_jiminy.common.utils.spaces.clip(data, space)[source]

Clip data from gym.Space to make sure it is within bounds.

Parameters:

  • data (Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray) – Data to clip.

  • space (Space[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray]) – gym.Space on which to operate.

Return type:

Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray

gym_jiminy.common.utils.spaces.contains(data, space, tol_abs=0.0, tol_rel=0.0)[source]

Check if all leaves of a nested data structure are within bounds of their respective gym.Space, up to some tolerance threshold. If both absolute and relative tolerances are provided, then satisfying only one of the two criteria is considered sufficient.

By design, it is always True for all spaces but gym.spaces.Box, gym.spaces.Discrete and gym.spaces.MultiDiscrete.

Parameters:

  • data (Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray) – Data structure to check.

  • space (Space[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray]) – gym.Space on which to operate.

  • tol_abs (float) – Absolute tolerance.

  • tol_rel (float) – Relative tolerance.

Return type:

bool

gym_jiminy.common.utils.spaces.build_reduce(fn, op, dataset, space, arity, *args, initializer=None, forward_bounds=True)[source]

Generate specialized callable applying transform and reduction on all leaves of given nested space.

Note

Original ordering of the leaves is preserved. More precisely, both transform and reduction will be applied recursively in keys order.

Warning

It is assumed without checking that all nested data structures are consistent together and with the space if provided. It holds true both data known at generation-time or runtime. Yet, it is only required for data provided at runtime if any to include the original data structure, so it may contain additional branches which will be ignored.

Warning

Providing additional data at runtime is supported but impede performance. Arity larger than 1 is not supported because the code path could not be fully specialized, causing dramatic slowdown.

Warning

There is no built-in ‘short-circuit’ mechanism, which means that it will go through all leaves systematically unless the reduction operator itself raises an exception.

Parameters:

  • fn (Callable[[...], ValueInT]) – Transform applied to every leaves of the nested data structures before performing the actual reduction. This function can perform in-place or out-of-place operations without restriction. None is not supported because it would be irrelevant. Note that if tracking the hierarchy during reduction is not necessary, then it would be way more efficient to first flatten the pre-allocated nested data structure once for all, and then perform reduction on this flattened view using the standard ‘functools.reduce’ method. Still, flattening at runtime using ‘flatten’ would still much slower than a specialized nested reduction.

  • op (Callable[[ValueOutT, ValueInT], ValueOutT] | None) – Optional reduction operator applied cumulatively on all leaves after transform. See ‘functools.reduce’ documentation for details. None to only apply transform on all leaves without reduction. This is useful when apply in-place transform.

  • data – Pre-allocated nested data structure. Optional if the space is provided but hardly relevant.

  • space (Space[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray] | None) – Container space on which to operate (eg gym.spaces.Dict or gym.spaces.Tuple). Optional iif the nested data structure is provided.

  • arity (Literal[0, 1] | None) – Arity of the generated callable. None to indicate that it must be determined at runtime, which is slower.

  • args (Any) – Extra arguments to systematically forward as transform input for all leaves. Note that, as for Python built-ins methods, keywords are not supported for the sake of efficiency.

  • initializer (Callable[[], ValueOutT] | None) – Function used to compute the initial value before starting reduction. Optional if the reduction operator has same input and output types. If None, then the value corresponding to the first leaf after transform will be used instead.

  • forward_bounds (bool) – Whether to forward the lower and upper bounds of the gym.Space associated with each leaf as transform input. In this case, they will be added after the data structure provided at runtime but before other extra arguments if any. It is up to the user to make sure all leaves have bounds, otherwise it will raise an exception at generation-time. This argument is ignored if not space is specified.

  • dataset (Sequence[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray])

Returns:

Fully-specialized reduction callable.

Return type:

Callable[[…], ValueOutT]

gym_jiminy.common.utils.spaces.build_map(fn, data, space, arity, *args, forward_bounds=True)[source]

Generate specialized callable returning applying out-of-place transform to all leaves of given nested space.

Warning

This method systematically allocates memory to store the resulting nested data structure, which is costly. If pre-allocation is possible, it would more efficient to use build_reduce without operator instead.

Warning

Providing additional data at runtime is supported but impede performance. Arity larger than 1 is not supported because the code path could not be fully specialized, causing dramatic slowdown.

Warning

It is assumed without check that all nested data structures are consistent together and with the space if provided. It holds true both data known at generation-time or runtime. Yet, it is only required for data provided at runtime if any to include the original data structure, so it may contain additional branches which will be ignored.

Parameters:

  • fn (Callable[[...], ValueT]) – Transform applied to every leaves of the nested data structures. This function is supposed to allocate its own memory while performing some out-of-place operations then return the outcome.

  • data (Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray | None) – Pre-allocated nested data structure. Optional iif the space is provided. This enables generating specialized random sampling methods for instance.

  • space (Space[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ndarray] | None) – gym.spaces.Dict on which to operate. Optional iif the nested data structure is provided.

  • arity (Literal[0, 1] | None) – Arity of the generated callable. None to indicate that it must be determined at runtime, which is slower.

  • args (Any) – Extra arguments to systematically forward as transform input for all leaves. Note that, as for Python built-ins methods, keywords are not supported for the sake of efficiency.

  • forward_bounds (bool) – Whether to forward the lower and upper bounds of the gym.Space associated with each leaf as transform input. In this case, they will be added after the data structure provided at runtime but before other extra arguments if any. It is up to the user to make sure all leaves have bounds, otherwise it will raise an exception at generation-time. This argument is ignored if not space is specified. Optional: True by default.

Returns:

Fully-specialized mapping callable.

Return type:

Callable[[], Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT]

gym_jiminy.common.utils.spaces.build_copyto(dst)[source]

Generate specialized copyto method for a given pre-allocated destination.

Parameters:

dst (Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray) – Nested data structure to be updated.

Return type:

Callable[[Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray], None]

gym_jiminy.common.utils.spaces.build_clip(data, space)[source]

Generate specialized clip method for some pre-allocated nested data structure and corresponding space.

Parameters:
Return type:

Callable[[], Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray]

gym_jiminy.common.utils.spaces.build_contains(data, space, tol_abs=0.0, tol_rel=0.0)[source]

Generate specialized contains method for some pre-allocated nested data structure and corresponding space.

Parameters:
Return type:

Callable[[], bool]

gym_jiminy.common.utils.spaces.build_normalize(space, dst, src=None, *, is_reversed=False)[source]

Generate a normalization or de-normalization method specialized for a given pre-allocated destination.

Note

The generated method applies element-wise de-normalization to all elements of the leaf spaces having finite bounds. For those that does not, it simply copies the value from ‘src’ to ‘dst’.

Warning

This method requires all leaf spaces to have type gym.spaces.Box with dtype ‘np.floating’.

Parameters:

  • dst (Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray) – Nested data structure to updated.

  • space (Space[Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray]) – Original (de-normalized) gym.Space on which to operate.

  • src (Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray | None) – Normalized nested data if ‘is_reversed’ is True, original data (de-normalized) otherwise. None to pass it at runtime. Optional: None by default.

  • is_reversed (bool) – True to de-normalize, False to normalize.

Return type:

Callable[[…], None]

gym_jiminy.common.utils.spaces.build_flatten(data_nested, data_flat=None, *, is_reversed=None)[source]

Generate a flattening or un-flattening method specialized for some pre-allocated nested data.

Note

Multi-dimensional leaf spaces are supported. Values will be flattened in 1D vectors using ‘C’ order (row-major). It ignores the actual memory layout the leaves of ‘data_nested’ and they are not required to have the same dtype as ‘data_flat’.

Parameters:

  • data_nested (Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray) – Nested data structure.

  • data_flat (Mapping[str, Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | Iterable[Mapping[str, StructNested[ValueT]] | Iterable[StructNested[ValueT]] | ValueT] | ndarray | None) – Flat array consistent with the nested data structure. Optional iif is_reversed is True. Optional: None by default.

  • is_reversed (bool | None) – True to update ‘data_flat’ (flattening), ‘data_nested’ otherwise (un-flattening). Optional: True if ‘data_flat’ is specified, False otherwise.

Return type:

Callable[[…], None]