From 17f43029e9250bf685ed0a7b8d5c1ecf18d5675e Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Mon, 25 Nov 2024 13:57:46 +0100 Subject: [PATCH] [gym_jiminy/common] Add 'AdaptLayoutObservation' that generalizes 'FilterObservation'. (#835) * [gym_jiminy/common] Add 'AdaptLayoutObservation' that generalizes 'FilterObservation'. * [gym_jiminy/common] Make sure that the initial pd state is within bounds. * [misc] Move from 'toml' to 'tomlkit' to fix heterogeneous array support. * [misc] Update documentation. * [misc] Fix typing. --- .../gym_jiminy/common/bases/interfaces.py | 2 +- .../proportional_derivative_controller.py | 5 +- .../common/gym_jiminy/common/envs/generic.py | 8 +- .../common/gym_jiminy/common/utils/math.py | 3 +- .../gym_jiminy/common/utils/pipeline.py | 23 +- .../common/gym_jiminy/common/utils/spaces.py | 19 +- .../gym_jiminy/common/wrappers/__init__.py | 3 +- .../gym_jiminy/common/wrappers/flatten.py | 6 +- .../common/wrappers/observation_filter.py | 191 ------ .../common/wrappers/observation_layout.py | 547 ++++++++++++++++++ python/gym_jiminy/common/setup.py | 5 +- python/jiminy_py/setup.py | 3 +- python/jiminy_py/src/jiminy_py/robot.py | 7 +- python/jiminy_py/src/jiminy_py/simulator.py | 31 +- python/jiminy_py/src/jiminy_py/tree.py | 6 +- 15 files changed, 631 insertions(+), 228 deletions(-) delete mode 100644 python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_filter.py create mode 100644 python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_layout.py diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py b/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py index c11d5e67e..f73eb0105 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py @@ -170,7 +170,7 @@ def compute_reward(self, # Similarly, `gym.Env` must be last to make sure all the other initialization # methods are called first. class InterfaceJiminyEnv( - InterfaceObserver[Obs, EngineObsType], + InterfaceObserver[Obs, EngineObsType], # type: ignore[type-var] InterfaceController[Act, np.ndarray], gym.Env[Obs, Act], Generic[Obs, Act]): diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py index b718ccfbc..ca2dc25f6 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/proportional_derivative_controller.py @@ -13,7 +13,7 @@ EncoderSensor, array_copyto) from ..bases import BaseObs, InterfaceJiminyEnv, BaseControllerBlock -from ..utils import fill +from ..utils import zeros, fill # Name of the n-th position derivative @@ -392,6 +392,9 @@ def __init__(self, # Initialize the controller super().__init__(name, env, update_ratio) + # Make sure that the state is within bounds + self._command_state[:2] = zeros(self.state_space) + # References to command acceleration for fast access self._command_acceleration = self._command_state[2] diff --git a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py index 0c6066801..794e882e9 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py @@ -435,7 +435,7 @@ def _get_measurements_space(self) -> spaces.Dict: joint_type = jiminy.get_joint_type(joint) if joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: sensor_position_lower = - np.pi - sensor_position_upper = np.pi + sensor_position_upper = + np.pi else: try: motor = self.robot.motors[sensor.motor_index] @@ -1243,7 +1243,7 @@ def evaluate(self, Optional: `None` by default. If not specified, then a strongly random seed will be generated by gym. :param horizon: Horizon of the simulation, namely maximum number of - steps before termination. `None` to disable. + env steps before termination. `None` to disable. Optional: Disabled by default. :param enable_stats: Whether to print high-level statistics after the simulation. @@ -1274,13 +1274,13 @@ def evaluate(self, self._initialize_seed(seed) # Initialize the simulation - obs, info = self.derived.reset() + env = self.derived + obs, info = env.reset() action, reward, terminated, truncated = None, None, False, False # Run the simulation info_episode = [info] try: - env = self.derived while horizon is None or self.num_steps < horizon: action = policy_fn( obs, action, reward, terminated, truncated, info) diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py index af5f8e750..c68855fac 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py @@ -1031,6 +1031,8 @@ def swing_from_vector( components of quaternions (x, y, z, w) and columns are the N independent orientations. """ + # pylint: disable=possibly-used-before-assignment + # Extract individual tilt components v_x, v_y, v_z = v_a @@ -1047,7 +1049,6 @@ def swing_from_vector( for i, q_i in enumerate(q.T): swing_from_vector((v_x[i], v_y[i], v_z[i]), q_i) else: - # pylint: disable=possibly-used-before-assignment eps_thr = np.sqrt(TWIST_SWING_SINGULAR_THR) eps_x = -TWIST_SWING_SINGULAR_THR < v_x < TWIST_SWING_SINGULAR_THR eps_y = -TWIST_SWING_SINGULAR_THR < v_y < TWIST_SWING_SINGULAR_THR diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py b/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py index ce81d67d6..f0f5c6c83 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/pipeline.py @@ -18,7 +18,7 @@ TypedDict, Literal, overload, cast) import h5py -import toml +import tomlkit import numpy as np import gymnasium as gym @@ -224,11 +224,14 @@ def build_pipeline(env_config: EnvConfig, :param env_config: Configuration of the environment, as a dict of type `EnvConfig`. - :param layers_config: Configuration of the blocks, as a list. The list is ordered from the lowest level layer to the highest, each element corresponding to the configuration of a individual layer, as a dict of type `LayerConfig`. + :param root_path: Optional path used as root for loading reference + trajectories from relative path if any. It will raise + an exception if required but not provided. + Optional: `None` by default. """ # Define helper to sanitize composition configuration def sanitize_composition_config(composition_config: CompositionConfig, @@ -554,13 +557,23 @@ def load_pipeline(fullpath: Union[str, pathlib.Path] :param: Fullpath of the configuration file. """ + # Extract root path from configuration file fullpath = pathlib.Path(fullpath) root_path, file_ext = fullpath.parent, fullpath.suffix + + # Load configuration file with open(fullpath, 'r') as f: if file_ext == '.json': - return build_pipeline(**json.load(f), root_path=root_path) - if file_ext == '.toml': - return build_pipeline(**toml.load(f), root_path=root_path) + # Parse JSON configuration file + all_config = json.load(f) + elif file_ext == '.toml': + # Parse TOML configuration file + all_config = tomlkit.load(f).unwrap() + else: + raise ValueError(f"File extension '{file_ext}' not supported.") + + # Build pipeline + return build_pipeline(**all_config, root_path=root_path) raise ValueError("Only json and toml formats are supported.") diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py b/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py index 7c2e3e842..55efe4969 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/spaces.py @@ -11,8 +11,8 @@ from collections.abc import Mapping, MutableMapping, Sequence, MutableSequence from typing import ( Any, Dict, Optional, Union, Sequence as SequenceT, Tuple, Literal, - Mapping as MappingT, Iterable, SupportsFloat, TypeVar, Type, Callable, - no_type_check, cast) + Mapping as MappingT, SupportsFloat, TypeVar, Type, Callable, no_type_check, + overload) import numba as nb import numpy as np @@ -30,7 +30,7 @@ StructNested = Union[MappingT[str, 'StructNested[ValueT]'], - Iterable['StructNested[ValueT]'], + SequenceT['StructNested[ValueT]'], ValueT] FieldNested = StructNested[str] DataNested = StructNested[np.ndarray] @@ -210,13 +210,24 @@ def copyto(dst: DataNested, src: DataNested) -> None: array_copyto(data, value) +@overload def copy(data: DataNestedT) -> DataNestedT: + ... + + +@overload +def copy(data: gym.Space[DataNestedT]) -> gym.Space[DataNestedT]: + ... + + +def copy(data: Union[DataNestedT, gym.Space[DataNestedT]] + ) -> Union[DataNestedT, gym.Space[DataNestedT]]: """Shallow copy recursively 'data' from `gym.Space`, so that only leaves are still references. :param data: Hierarchical data structure to copy without allocation. """ - return cast(DataNestedT, tree.unflatten_as(data, tree.flatten(data))) + return tree.unflatten_as(data, tree.flatten(data)) @no_type_check diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py index 3d450525c..b5e105acb 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/__init__.py @@ -1,12 +1,13 @@ # pylint: disable=missing-module-docstring -from .observation_filter import FilterObservation +from .observation_layout import AdaptLayoutObservation, FilterObservation from .observation_stack import StackObservation from .normalize import NormalizeAction, NormalizeObservation from .flatten import FlattenAction, FlattenObservation __all__ = [ + 'AdaptLayoutObservation', 'FilterObservation', 'StackObservation', 'NormalizeObservation', diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/flatten.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/flatten.py index 4b48b03df..d3e231609 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/flatten.py +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/flatten.py @@ -44,10 +44,10 @@ def __init__(self, """ # Find most appropriate dtype if not specified if dtype is None: - obs_flat = tuple( - value.dtype for value in tree.flatten(env.observation)) if env.observation: - dtype = reduce(np.promote_types, obs_flat) + dtype_all = [ + value.dtype for value in tree.flatten(env.observation)] + dtype = reduce(np.promote_types, dtype_all) else: dtype = np.float64 diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_filter.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_filter.py deleted file mode 100644 index 2f1664d6e..000000000 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_filter.py +++ /dev/null @@ -1,191 +0,0 @@ -"""This module implements a block transformation for filtering out some of the -keys of the observation space of an environment that may be arbitrarily nested. -""" -from operator import getitem -from functools import reduce -from collections import OrderedDict -from typing import ( - Sequence, Set, Tuple, Union, Generic, TypeVar, Type, TypeAlias, - no_type_check) - -import gymnasium as gym -from jiminy_py.tree import ( - flatten_with_path, issubclass_mapping, issubclass_sequence) - -from ..bases import (NestedObs, - Act, - InterfaceJiminyEnv, - BaseTransformObservation) -from ..utils import DataNested, copy - - -SpaceOrData = TypeVar( - 'SpaceOrData', bound=Union[DataNested, gym.Space[DataNested]]) -FilteredObs: TypeAlias = NestedObs - - -@no_type_check -def _copy_filtered(data: SpaceOrData, - path_filtered_leaves: Set[Tuple[str, ...]]) -> SpaceOrData: - """Partially shallow copy some nested data structure, so that all leaves - being filtered in are still references but their corresponding containers - are copies. - - :param data: Hierarchical data structure to copy without allocation. - :param path_filtered_leaves: Set gathering the paths of all leaves that - must be kept. Each path is a tuple of keys - to access a given leaf recursively. - """ - # Special handling if no leaf to filter has been specified - if not path_filtered_leaves: - data_type = type(data) - if issubclass_mapping(data_type) or issubclass_sequence(data_type): - return data_type() - return data - - # Shallow copy the whole data structure - out = copy(data) - - # Convert all parent containers to mutable dictionary - type_filtered_nodes: Sequence[Type] = [] - out_flat = flatten_with_path(out) - for key_nested, _ in out_flat: - if key_nested not in path_filtered_leaves: - continue - for i in range(1, len(key_nested) + 1): - # Extract parent container - *key_nested_parent, key_leaf = key_nested[:i] - if key_nested_parent: - *key_nested_container, key_parent = key_nested_parent - container = reduce(getitem, key_nested_container, out) - parent = container[key_parent] - else: - parent = out - - # Convert parent container to mutable dictionary - parent_type = type(parent) - type_filtered_nodes.append(parent_type) - if parent_type in (list, dict, OrderedDict): - continue - if issubclass_mapping(parent_type): - parent = dict(parent.items()) - elif issubclass_sequence(parent_type): - parent = dict(enumerate(parent)) - else: - raise NotImplementedError( - f"Unsupported container type: '{parent_type}'") - - # Re-assign parent data structure - if key_nested_parent: - container[key_parent] = parent - else: - out = parent - - # Remove unnecessary keys - for key_nested, _ in out_flat: - if key_nested in path_filtered_leaves: - continue - for i in range(len(key_nested))[::-1]: - if any(key_nested[:i] == path[:i] - for path in path_filtered_leaves): - break - *key_nested_parent, key_leaf = key_nested[:(i + 1)] - try: - parent = reduce(getitem, key_nested_parent, out) - del parent[key_leaf] - except KeyError: - # Some nested keys may have been deleted previously - pass - - # Restore original parent container types - parent_type_it = iter(type_filtered_nodes[::-1]) - for key_nested, _ in out_flat[::-1]: - if key_nested not in path_filtered_leaves: - continue - for i in range(1, len(key_nested) + 1)[::-1]: - # Extract parent container - *key_nested_parent, _ = key_nested[:i] - if key_nested_parent: - *key_nested_container, key_parent = key_nested_parent - container = reduce(getitem, key_nested_container, out) - parent = container[key_parent] - else: - parent = out - - # Restore original container type if not already done - parent_type = next(parent_type_it) - if isinstance(parent, parent_type): - continue - if issubclass_mapping(parent_type): - parent = parent_type(tuple(parent.items())) - elif issubclass_sequence(parent_type): - parent = parent_type(tuple(parent.values())) - - # Re-assign output data structure - if key_nested_parent: - container[key_parent] = parent - else: - out = parent - return out - - -class FilterObservation( - BaseTransformObservation[FilteredObs, NestedObs, Act], - Generic[NestedObs, Act]): - """Filter nested observation space. - - This wrapper does nothing but providing an observation only exposing a - subset of all the leaves of the original observation space. For flattening - the observation space after filtering, you should wrap the environment with - `FlattenObservation` as yet another layer. - """ - def __init__(self, - env: InterfaceJiminyEnv[NestedObs, Act], - nested_filter_keys: Sequence[ - Union[Sequence[Union[str, int]], Union[str, int]]] - ) -> None: - # Make sure that the observation space derives from 'gym.spaces.Dict' - assert isinstance(env.observation_space, gym.spaces.Dict) - - # Make sure all nested keys are stored in sequence - assert not isinstance(nested_filter_keys, str) - self.nested_filter_keys: Sequence[Sequence[Union[str, int]]] = [] - for key_nested in nested_filter_keys: - if isinstance(key_nested, (str, int)): - key_nested = (key_nested,) - self.nested_filter_keys.append(tuple(key_nested)) - - # Get all paths associated with leaf values that must be stacked - self.path_filtered_leaves: Set[Tuple[Union[str, int], ...]] = set() - for path, _ in flatten_with_path(env.observation_space): - if any(path[:len(e)] == e for e in self.nested_filter_keys): - self.path_filtered_leaves.add(path) - - # Make sure that some keys are preserved - if not self.path_filtered_leaves: - raise ValueError( - "At least one observation leaf must be preserved.") - - # Initialize base class - super().__init__(env) - - # Bind observation of the environment for all filtered keys. - # Note that all parent containers of the filtered leaves must be - # constructible from standard `dict` or `tuple` objects, which is the - # case for all standard `gym.Space`. - self.observation = _copy_filtered( - self.env.observation, self.path_filtered_leaves) - - def _initialize_observation_space(self) -> None: - """Configure the observation space. - - It gathers a subset of all the leaves of the original observation space - without any further processing. - """ - self.observation_space = _copy_filtered( - self.env.observation_space, self.path_filtered_leaves) - - def transform_observation(self) -> None: - """No-op transform since the transform observation is sharing memory - with the wrapped one since it is just a partial view. - """ diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_layout.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_layout.py new file mode 100644 index 000000000..c69780d7f --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_layout.py @@ -0,0 +1,547 @@ +# mypy: disable-error-code="arg-type, index, assignment, union-attr" +"""This module implements a block transformation for adapt the layout of the +observation space of an environment that may be arbitrarily nested. +""" +from operator import getitem +from functools import reduce +from collections import OrderedDict +from typing import ( + Sequence, Tuple, List, Union, Generic, TypeVar, Type, Optional, Dict, cast, + overload) +from typing_extensions import Unpack + +import numpy as np +import gymnasium as gym +from jiminy_py.tree import ( + flatten_with_path, issubclass_mapping, issubclass_sequence) + +from ..bases import (Act, + NestedObs, + InterfaceJiminyEnv, + BaseTransformObservation) +from ..utils import DataNested, copy + + +MaybeNestedObs = TypeVar('MaybeNestedObs', bound=DataNested) +OtherMaybeNestedObs = TypeVar('OtherMaybeNestedObs', bound=DataNested) + +Key = Union[str, int] +NestedKey = Tuple[Key, ...] +Slice = Union[Tuple[()], int, Tuple[Optional[int], Optional[int]]] +ArrayBlockSpec = Sequence[Slice] +NestedData = Union[NestedKey, Tuple[Unpack[NestedKey], ArrayBlockSpec]] + + +@overload +def _adapt_layout(data: DataNested, + layout: Sequence[Tuple[NestedKey, NestedData]] + ) -> DataNested: + ... + + +@overload +def _adapt_layout(data: gym.Space[DataNested], + layout: Sequence[Tuple[NestedKey, NestedData]] + ) -> gym.Space[DataNested]: + ... + + +def _adapt_layout( + data: Union[DataNested, gym.Space[DataNested]], + layout: Sequence[Tuple[NestedKey, NestedData]] + ) -> Union[DataNested, gym.Space[DataNested]]: + """Extract subtrees and leaves from a given input nested data structure and + rearrange them an arbitrary output nested data structure according to some + prescribed layout. + + All the leaves are still sharing memomy with their original counterpart as + they are references to the same objects. However, all the containers are + completely independent, even when extracting a complete subtree. + + .. note:: + This method can be used to filter out subtrees of a given nested data + structure by partially shallow copying it. + + .. warning:: + All the output containers are freshly created rather than copies. This + means that they must be constructible from plain-old python containers + (dict, list, tuple, ...) without any extra arguments. Any custom + attributes dynamically added would be lost if any. This is typically + the case for all standard `gym.Space`. + + :param data: Hierarchical data structure from which to extract data without + memory allocation for leaves. + :param layout: Sequence of tuples `(nested_key_out, nested_key_in)` mapping + the desired path in the output data structure to the + original path in input data structure. These tuples are + guaranteed to be processed in order. The same path may + appear multiple time in the output data structure. If so, + the corresponding subtrees while be aggregated in sequence. + If the parent node of the first subtree being considered is + already a sequence, then subsequent extracted subtrees will + be appended to it directly. If not, then a dedicated + top-level sequence container while be created first. + """ + # We need to guarantee that all containers of the output data structure are + # mutable while building it, and only at the end to convert them back to + # their respective desired type. + # Since the output data structure is not known in advance, the only option + # is to keep track of the desired container type while building the output + # data structure. + + # Determine if data is a gym.Space or a "classical" nested data structure + is_space = isinstance(data, gym.Space) + + # Shallow copy the input data structure + data = copy(data) + + # Convert all parent containers to their mutable counterpart. + container_types_in: Dict[NestedKey, Type] = {} + for key_nested, _ in flatten_with_path(data): + for i in range(1, len(key_nested) + 1): + # Extract parent container + key_nested_parent = key_nested[:(i - 1)] + if key_nested_parent: + *key_nested_container, key_parent = key_nested_parent + container = reduce(getitem, key_nested_container, data) + parent = container[key_parent] + else: + parent = data + + # Convert parent container to mutable dictionary + parent_type = type(parent) + if key_nested_parent not in container_types_in: + container_types_in[key_nested_parent] = parent_type + if parent_type in (list, dict, OrderedDict): + continue + if issubclass_mapping(parent_type): + parent = dict(parent.items()) + elif issubclass_sequence(parent_type): + parent = list(parent) + else: + raise NotImplementedError( + f"Unsupported container type: '{parent_type}'") + + # Re-assign parent data structure + if key_nested_parent: + container[key_parent] = parent + else: + data = parent + + # Build the output data structure sequentially + container_types_out: Dict[NestedKey, Type] = {} + out: Optional[DataNested] = None + for nested_key_out, nested_spec_in in layout: + # Make sure that requested nested data is a subtree + if not nested_spec_in: + raise ValueError("Input nested keys must not be empty.") + + # Split nested keys from block specification if any + block_spec_in: Optional[ArrayBlockSpec] = None + if isinstance(nested_spec_in[-1], (tuple, list)): + nested_key_in = nested_spec_in[:-1] + block_spec_in = nested_spec_in[-1] + else: + nested_key_in = cast(NestedKey, nested_spec_in) + + # Extract the input chunk recursively + value_in = reduce(getitem, nested_key_in, data) + if block_spec_in is not None: + # Convert array block specification to slices + slices = [] + for start_end in block_spec_in: + if isinstance(start_end, int): + slices.append(start_end) + elif not start_end: + slices.append(slice(None,)) + else: + slices.append(slice(*start_end)) + slices = tuple(slices) + + # Extract sub-space or array view depending on the input + if is_space: + assert isinstance(value_in, gym.spaces.Box) + assert isinstance(value_in.dtype, np.dtype) and issubclass( + value_in.dtype.type, (np.floating, np.integer)) + low, high = value_in.low[slices], value_in.high[slices] + value_in = gym.spaces.Box(low=low, + high=high, + shape=low.shape, + dtype=value_in.dtype.type) + else: + assert isinstance(value_in, np.ndarray) + value_in = value_in[slices] + + # Deal with the special case where the output nested key is empty + if not nested_key_out: + if out is None: + # Promote the extracted value as output + out = value_in + else: + # Encapsulate the output in sequence container if it was not + # the case originally. The whole hierarchy of container types + # must be shifted one level deeper, while the new top-level + # container type is the default tuple. + if not isinstance(out, list): + out = [out,] + for path, type_ in tuple(container_types_out.items()): + container_types_out[(0, *path)] = type_ + del container_types_out[path] + if isinstance(data, gym.Space): + container_types_out[()] = gym.spaces.Tuple + else: + container_types_out[()] = tuple + + # Update out nested key to account for index in sequence + nested_key_out = (len(out),) + out.append(value_in) + + # Add extracted input chunks to output + value_out = out + depth = len(nested_key_out) + for i in range(depth): + # Extract key and subkey + key = nested_key_out[i] + subkey = nested_key_out[i + 1] if i + 1 < depth else None + + if isinstance(key, str): + # Initialize the out container is not done before + if out is None: + value_out = out = {} + if is_space: + container_types_out[()] = gym.spaces.Dict + else: + container_types_out[()] = dict + assert isinstance(value_out, dict) + + # Create new branch if not the final key in path, otherwise + # add extracted value as leaf. Then, extract child node. + if key not in value_out or subkey is not None: + if key not in value_out: + if subkey is None: + value_out[key] = value_in + else: + value_out[key] = {} + path = nested_key_out[:(i + 1)] + if is_space: + container_types_out[path] = gym.spaces.Dict + else: + container_types_out[path] = dict + value_out = value_out[key] + continue + elif isinstance(key, int): + # Initialize the out container is not done before + if out is None: + value_out = out = [] + if is_space: + container_types_out[()] = gym.spaces.Tuple + else: + container_types_out[()] = tuple + assert isinstance(value_out, list) + + # Just extract child node if not the final key in path + if subkey is not None: + value_out = value_out[key] + continue + + # The final node expected to be sequence container. Encapsulate the + # output in sequence container if not already the case, while + # adapting the hierarchy of output container types accordingly. + # Then append the extract value. + assert value_out is not None + if not isinstance(value_out[key], list): + value_out[key] = [value_out[key],] + for path, type_ in tuple(container_types_out.items()): + root_path, child_path = path[:depth], path[depth:] + if root_path == nested_key_out: + path_ = (*nested_key_out, 0, *child_path) + container_types_out[path_] = type_ + del container_types_out[path] + if is_space: + container_types_out[nested_key_out] = gym.spaces.Tuple + else: + container_types_out[nested_key_out] = tuple + nested_key_out = (*nested_key_out, len(value_out[key])) + value_out[key].append(value_in) + + # Extract copied out container types + for path, type_ in container_types_in.items(): + root_path = path[:len(nested_key_in)] + child_path = path[len(nested_key_in):] + if root_path == nested_key_in: + container_types_out[(*nested_key_out, *child_path)] = type_ + + # Restore original parent container types + path_all, _ = zip(*flatten_with_path(out)) + depth = max(map(len, path_all)) + for i in range(depth)[::-1]: + for key_nested in path_all: + # Skip if the node is not a container + if len(key_nested) <= i: + continue + + # Extract parent container + key_nested_parent = key_nested[:i] + if key_nested_parent: + *key_nested_container, key_parent = key_nested_parent + container = reduce(getitem, key_nested_container, out) + parent = container[key_parent] + else: + parent = out + + # Restore original container type if not already done + parent_type = container_types_out[key_nested_parent] + if isinstance(parent, parent_type): + continue + if issubclass_mapping(parent_type): + parent = parent_type(tuple(parent.items())) + elif issubclass_sequence(parent_type): + parent = parent_type(tuple(parent)) + + # Re-assign output data structure + if key_nested_parent: + container[key_parent] = parent + else: + out = parent + + assert out is not None + return out + + +class AdaptLayoutObservation( + BaseTransformObservation[OtherMaybeNestedObs, MaybeNestedObs, Act], + Generic[OtherMaybeNestedObs, MaybeNestedObs, Act]): + """Adapt the data structure of the original nested observation space, + by filtering out some leaves and/or re-ordering others. + + This wrapper does nothing but exposing a subset of all the leaves of the + original observation space with a completely independent data structure. + For flattening the observation space after filtering, you should wrap the + environment with `FlattenObservation` as yet another layer. + + It is possible to operate on subtrees directly without going all the way + down each leaf. Similarly, one can extract slices (possibly not contiguous) + of `gym.spaces.Box` spaces. Extra leaves can be added to subtrees of + original data that has already been re-mapped, knowing that the items of + the layout are always processed in order. Moreover, the same output key can + appear multiple times. In such case, all the associated values will be + stacked as a tuple while preserving their order. + + Let us consider the following nested observation space: + + gym.spaces.Dict( + x1=gym.spaces.Tuple([ + gym.spaces.Box(float('-inf'), float('inf'), (7, 6, 5, 4)), + ]), + x2=gym.spaces.Dict( + y1=gym.spaces.Dict( + z1=gym.spaces.Discrete(5) + ), + y2=gym.spaces.Tuple([ + gym.spaces.Discrete(2), + gym.spaces.Box(float('-inf'), float('inf'), (2, 3)), + gym.spaces.Discrete(3), + ]))) + + Here is an example that aggregates one leaf to a leaf of a subtree that has + already been remapped: + + [((), ("x2", "y2")), ((0,), ("x1",))] + + gym.spaces.Tuple([ + gym.spaces.Tuple([ + gym.spaces.Discrete(2), + gym.spaces.Tuple([ + gym.spaces.Box(float('-inf'), float('inf'), (7, 6, 5, 4)), + ]) + ]), + gym.spaces.Box(float('-inf'), float('inf'), (2, 3)), + gym.spaces.Discrete(3), + ])) + + Here is an example that aggregate one leaf and one multi-dimensional slice + of a Box space (array[:, 1:3]) under two separate keys of a nested dict: + + [(("A", "B1"), ("x2", "y2", 1)), + (("A", "B2"), ("x1", 0, [(), (1, 4), (1, 3)]))] + + gym.spaces.Dict( + A=gym.spaces.Dict( + B1=gym.spaces.Box(float('-inf'), float('inf'), (2, 3)), + B2=gym.spaces.Box(float('-inf'), float('inf'), (7, 3, 2, 4)) + )) + + Here is an example that aggregate two subtree in the same nested structure: + + [((), ("x2",)), (("y1", "z2"), ("x1",))] + + gym.spaces.Dict( + y1=gym.spaces.Dict( + z1=gym.spaces.Discrete(5) + z2=gym.spaces.Tuple([ + gym.spaces.Box(float('-inf'), float('inf'), (7, 6, 5, 4)), + ]) + ), + y2=gym.spaces.Tuple([ + gym.spaces.Discrete(2), + gym.spaces.Box(float('-inf'), float('inf'), (2, 3)), + gym.spaces.Discrete(3), + ])) + """ + def __init__(self, + env: InterfaceJiminyEnv[MaybeNestedObs, Act], + layout: Sequence[Tuple[ + Union[NestedKey, Key], NestedData]]) -> None: + """ + :param env: Base or already wrapped jiminy environment. + :param layout: Sequence of tuples `(nested_key_out, nested_key_in)` + mapping the desired path in the output data structure to + the original path in input data structure. These tuples + are guaranteed to be processed in order. The same path + may appear multiple time in the output data structure. + If so, the corresponding subtrees while be aggregated in + sequence. If the parent node of the first subtree being + considered is already a sequence, then all subsequent + extracted subtrees will be appended to it directly. If + not, then a dedicated top-level sequence container while + be created first. + """ + # Make sure that some keys are preserved + if not layout: + raise ValueError( + "The resulting observation space must not be empty.") + + # Backup user-specified layout while making sure all nested keys are + # stored in sequence. + self.layout: Sequence[Tuple[NestedKey, NestedData]] = [] + for nested_key_out, nested_spec_in in layout: + if isinstance(nested_key_out, (str, int)): + nested_key_out = (nested_key_out,) + self.layout.append((nested_key_out, nested_spec_in)) + + # Initialize base class + super().__init__(env) + + # Bind observation of the environment for all extracted keys + self.observation = _adapt_layout(self.env.observation, self.layout) + + def _initialize_observation_space(self) -> None: + """Configure the observation space. + + It gathers a subset of all the leaves of the original observation space + without any further processing. + """ + self.observation_space = _adapt_layout( + self.env.observation_space, self.layout) + + def transform_observation(self) -> None: + """No-op transform since the transform observation is sharing memory + with the wrapped one since it is just a partial view. + """ + + +class FilterObservation( + BaseTransformObservation[NestedObs, NestedObs, Act], + Generic[NestedObs, Act]): + """Filter nested observation space. + + This wrapper does nothing but providing an observation only exposing a + subset of all the leaves of the original observation space. This wrapper is + a specialization of `AdaptLayoutObservation`, which is more generic. For + flattening the observation space after filtering, you should wrap the + environment with `FlattenObservation` as yet another layer. + + Note that it is possible to operate on subtrees directly without going all + the way down each leaf. + + Beware that the original ordered of leaves within their parent container is + maintain, whatever the order in which they appear in the specified list of + filtered nested keys. + + Let us consider the following nested observation space: + + gym.spaces.Dict( + x1=gym.spaces.Tuple([ + gym.spaces.Box(float('-inf'), float('inf'), (7, 6, 5, 4)), + ]), + x2=gym.spaces.Dict( + y1=gym.spaces.Dict( + z1=gym.spaces.Discrete(5) + ), + y2=gym.spaces.Tuple([ + gym.spaces.Discrete(2), + gym.spaces.Box(float('-inf'), float('inf'), (2, 3)), + gym.spaces.Discrete(3), + ]))) + + Here is an example that filter out part of a sequence container and a + mapping container: + + [("x2", "y2", 2), ("x2", "y2", 0), ("x1",)] + + gym.spaces.Dict( + x1=gym.spaces.Tuple([ + gym.spaces.Box(float('-inf'), float('inf'), (7, 6, 5, 4)), + ]), + x2=gym.spaces.Dict( + y2=gym.spaces.Tuple([ + gym.spaces.Discrete(2), + gym.spaces.Discrete(3), + ]))) + """ + def __init__(self, + env: InterfaceJiminyEnv[NestedObs, Act], + nested_filter_keys: Sequence[Union[NestedKey, Key]]) -> None: + # Make sure that the top-most observation space is a container + space_cls = type(env.observation_space) + assert issubclass_mapping(space_cls) or issubclass_sequence(space_cls) + + # Make sure all nested keys are stored in sequence + assert not isinstance(nested_filter_keys, str) + self.nested_filter_keys: Sequence[NestedKey] = [] + for key_nested in nested_filter_keys: + if isinstance(key_nested, (str, int)): + key_nested = (key_nested,) + self.nested_filter_keys.append(tuple(key_nested)) + + # Get all paths associated with leaf values that must be kept. + # Re-order filtered leaves to match the original nested data structure. + path_filtered_leaves: List[NestedKey] = [] + for path, _ in flatten_with_path(env.observation_space): + if any(path[:len(e)] == e for e in self.nested_filter_keys): + if path not in path_filtered_leaves: + path_filtered_leaves.append(path) + + # Backup the layout mapping + self._layout: Sequence[Tuple[NestedKey, NestedData]] = [] + for nested_key_in in path_filtered_leaves: + if isinstance(nested_key_in[-1], int): + nested_key_out = nested_key_in[:-1] + else: + nested_key_out = nested_key_in + self._layout.append((nested_key_out, nested_key_in)) + + # Make sure that some keys are preserved + if not path_filtered_leaves: + raise ValueError( + "At least one observation leaf must be preserved.") + + # Initialize base class + super().__init__(env) + + # Bind observation of the environment for all filtered keys + self.observation = _adapt_layout(self.env.observation, self._layout) + + def _initialize_observation_space(self) -> None: + """Configure the observation space. + + It gathers a subset of all the leaves of the original observation space + without any further processing. + """ + self.observation_space = _adapt_layout( + self.env.observation_space, self._layout) + + def transform_observation(self) -> None: + """No-op transform since the transform observation is sharing memory + with the wrapped one since it is just a partial view. + """ diff --git a/python/gym_jiminy/common/setup.py b/python/gym_jiminy/common/setup.py index 019e99d8f..c6bc900d0 100644 --- a/python/gym_jiminy/common/setup.py +++ b/python/gym_jiminy/common/setup.py @@ -57,7 +57,10 @@ # - `gym` has been replaced by `gymnasium` for 0.26.0+ # - 0.28.0: fully typed # - bound version for resilience to recurrent API breakage - "gymnasium>=0.28,<1.0" + "gymnasium>=0.28,<1.0", + # For backward compatibility of latest Python typing features + # - 'Unpack' was introduced with Python 3.11 + "typing_extensions>=4.1.0", ], extras_require=extras, zip_safe=False diff --git a/python/jiminy_py/setup.py b/python/jiminy_py/setup.py index b6f633261..1527f0f96 100644 --- a/python/jiminy_py/setup.py +++ b/python/jiminy_py/setup.py @@ -99,7 +99,7 @@ def finalize_options(self) -> None: # see: https://numpy.org/devdocs/dev/depending_on_numpy.html np_req, # Parser for Jiminy's hardware description file. - "toml", + "tomlkit", # Standalone cross-platform mesh visualizer used as Viewer's backend. # Panda3d is NOT supported by PyPy even if built from source. # - 1.10.12 fixes numerous bugs @@ -154,7 +154,6 @@ def finalize_options(self) -> None: # Stub for static type checking "types-psutil>=5.9.5.20240511", "types-Pillow", - "types-toml", "types-tqdm", # Check PEP8 conformance of Python native code "flake8", diff --git a/python/jiminy_py/src/jiminy_py/robot.py b/python/jiminy_py/src/jiminy_py/robot.py index d07273889..823898d78 100644 --- a/python/jiminy_py/src/jiminy_py/robot.py +++ b/python/jiminy_py/src/jiminy_py/robot.py @@ -13,7 +13,7 @@ from types import ModuleType from typing import Optional, Dict, Any, Sequence, Literal, Set, List, get_args -import toml +import tomlkit import numpy as np import trimesh import trimesh.parent @@ -513,7 +513,7 @@ def generate_default_hardware_description_file( hardware_path = str(pathlib.Path( urdf_path).with_suffix('')) + '_hardware.toml' with open(hardware_path, 'w') as f: - toml.dump(hardware_info, f) + tomlkit.dump(hardware_info, f) def load_hardware_description_file( @@ -546,7 +546,8 @@ def load_hardware_description_file( else: LOGGER.setLevel(logging.ERROR) - hardware_info = toml.load(hardware_path) + with open(hardware_path, 'r') as f: + hardware_info = tomlkit.load(f).unwrap() extra_info = hardware_info.pop('Global', {}) motors_info = hardware_info.pop('Motor', {}) sensors_info = hardware_info.pop('Sensor', {}) diff --git a/python/jiminy_py/src/jiminy_py/simulator.py b/python/jiminy_py/src/jiminy_py/simulator.py index b2e8c104c..717781d47 100644 --- a/python/jiminy_py/src/jiminy_py/simulator.py +++ b/python/jiminy_py/src/jiminy_py/simulator.py @@ -15,10 +15,10 @@ from functools import partial from typing import Any, List, Dict, Optional, Union, Sequence, Callable -import toml +import tomlkit import numpy as np -from . import core as jiminy +from . import core as jiminy, tree from .robot import BaseJiminyRobot, generate_default_hardware_description_file from .dynamics import Trajectory from .log import UpdateHook, read_log, build_robot_from_log @@ -969,10 +969,18 @@ def export_options(self, config_path: Union[str, os.PathLike]) -> None: generated file. The extension '.toml' will be enforced. """ + # Get all simulation options + simu_options = self.get_simulation_options() + + # Convert all numpy array options to list + simu_options = tree.unflatten_as(simu_options, [ + value.tolist() if isinstance(value, np.ndarray) else value + for path, value in tree.flatten_with_path(simu_options)]) + + # Dump all simulation options in the same configuration file config_path = pathlib.Path(config_path).with_suffix('.toml') with open(config_path, 'w') as f: - toml.dump( - self.get_simulation_options(), f, toml.TomlNumpyEncoder()) + tomlkit.dump(simu_options, f) # type: ignore[arg-type] def import_options(self, config_path: Union[str, os.PathLike]) -> None: """Import all the options of the simulator at once, ie the engine @@ -985,7 +993,7 @@ def import_options(self, config_path: Union[str, os.PathLike]) -> None: :param config_path: Full path of the configuration file to load. """ def deep_update(original: Dict[str, Any], - new_dict: Dict[str, Any], + new_dict: Union[Dict[str, Any], tomlkit.TOMLDocument], *, _key_root: str = "") -> Dict[str, Any]: """Updates `original` dict with values from `new_dict` recursively. If a new key should be introduced, then an error is thrown instead. @@ -1009,6 +1017,13 @@ def deep_update(original: Dict[str, Any], original[key] = new_dict[key] return original - options = deep_update( - self.get_simulation_options(), toml.load(str(config_path))) - self.set_simulation_options(options) + # Load (partial) simulation options + with open(config_path, 'r') as f: + simu_options = tomlkit.load(f).unwrap() + + # Fill any missing key with their current value + simu_options_full = deep_update( + self.get_simulation_options(), simu_options) + + # Set all options at once + self.set_simulation_options(simu_options_full) diff --git a/python/jiminy_py/src/jiminy_py/tree.py b/python/jiminy_py/src/jiminy_py/tree.py index c3dd75f59..4c7e4200b 100644 --- a/python/jiminy_py/src/jiminy_py/tree.py +++ b/python/jiminy_py/src/jiminy_py/tree.py @@ -20,13 +20,13 @@ from itertools import chain, starmap from collections.abc import (Mapping, ValuesView, Sequence, Set) from typing import ( - Any, Union, Mapping as MappingT, Iterable, Iterator as Iterator, Tuple, - TypeVar, Callable, Type) + Any, Union, Mapping as MappingT, Sequence as SequenceT, Iterable, + Iterator as Iterator, Tuple, TypeVar, Callable, Type) ValueT = TypeVar('ValueT') StructNested = Union[MappingT[str, 'StructNested[ValueT]'], - Iterable['StructNested[ValueT]'], + SequenceT['StructNested[ValueT]'], ValueT]