-from typing import Tuple, Union, Callable, List, Dict, Any, Optional
-
-import copy
-import importlib
-import numpy as np
-from collections import defaultdict
-
-from collections.abc import Mapping, MutableMapping
-
-from fancy_gym.utils.make_env_helpers import make_bb
-from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
-
-from gymnasium import register as gym_register
-from gymnasium import make as gym_make
-from gymnasium.envs.registration import registry as gym_registry
-
-
-class DefaultMPWrapper(RawInterfaceWrapper):
- @property
- def context_mask(self):
- """
- Returns boolean mask of the same shape as the observation space.
- It determines whether the observation is returned for the contextual case or not.
- This effectively allows to filter unwanted or unnecessary observations from the full step-based case.
- E.g. Velocities starting at 0 are only changing after the first action. Given we only receive the
- context/part of the first observation, the velocities are not necessary in the observation for the task.
- Returns:
- bool array representing the indices of the observations
- """
- # If the env already defines a context_mask, we will use that
- if hasattr(self.env, 'context_mask'):
- return self.env.context_mask
-
- # Otherwise we will use the whole observation as the context. (Write a custom MPWrapper to change this behavior)
- return np.full(self.env.observation_space.shape, True)
-
- @property
- def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
- """
- Returns the current position of the action/control dimension.
- The dimensionality has to match the action/control dimension.
- This is not required when exclusively using velocity control,
- it should, however, be implemented regardless.
- E.g. The joint positions that are directly or indirectly controlled by the action.
- """
- assert hasattr(self.env, 'current_pos'), 'DefaultMPWrapper was unable to access env.current_pos. Please write a custom MPWrapper (recommended) or expose this attribute directly.'
- return self.env.current_pos
-
- @property
- def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
- """
- Returns the current velocity of the action/control dimension.
- The dimensionality has to match the action/control dimension.
- This is not required when exclusively using position control,
- it should, however, be implemented regardless.
- E.g. The joint velocities that are directly or indirectly controlled by the action.
- """
- assert hasattr(self.env, 'current_vel'), 'DefaultMPWrapper was unable to access env.current_vel. Please write a custom MPWrapper (recommended) or expose this attribute directly.'
- return self.env.current_vel
-
-
-_BB_DEFAULTS = {
- 'ProMP': {
- 'wrappers': [],
- 'trajectory_generator_kwargs': {
- 'trajectory_generator_type': 'promp'
- },
- 'phase_generator_kwargs': {
- 'phase_generator_type': 'linear'
- },
- 'controller_kwargs': {
- 'controller_type': 'motor',
- 'p_gains': 1.0,
- 'd_gains': 0.1,
- },
- 'basis_generator_kwargs': {
- 'basis_generator_type': 'zero_rbf',
- 'num_basis': 5,
- 'num_basis_zero_start': 1,
- 'basis_bandwidth_factor': 3.0,
- },
- 'black_box_kwargs': {
- }
- },
- 'DMP': {
- 'wrappers': [],
- 'trajectory_generator_kwargs': {
- 'trajectory_generator_type': 'dmp'
- },
- 'phase_generator_kwargs': {
- 'phase_generator_type': 'exp'
- },
- 'controller_kwargs': {
- 'controller_type': 'motor',
- 'p_gains': 1.0,
- 'd_gains': 0.1,
- },
- 'basis_generator_kwargs': {
- 'basis_generator_type': 'rbf',
- 'num_basis': 5
- },
- 'black_box_kwargs': {
- }
- },
- 'ProDMP': {
- 'wrappers': [],
- 'trajectory_generator_kwargs': {
- 'trajectory_generator_type': 'prodmp',
- 'duration': 2.0,
- 'weights_scale': 1.0,
- },
- 'phase_generator_kwargs': {
- 'phase_generator_type': 'exp',
- 'tau': 1.5,
- },
- 'controller_kwargs': {
- 'controller_type': 'motor',
- 'p_gains': 1.0,
- 'd_gains': 0.1,
- },
- 'basis_generator_kwargs': {
- 'basis_generator_type': 'prodmp',
- 'alpha': 10,
- 'num_basis': 5,
- },
- 'black_box_kwargs': {
- }
- }
-}
-
-KNOWN_MPS = list(_BB_DEFAULTS.keys())
-_KNOWN_MPS_PLUS_ALL = KNOWN_MPS + ['all']
-ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS = {mp_type: [] for mp_type in _KNOWN_MPS_PLUS_ALL}
-MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS = {}
-
-
-[docs]def register(
-
id: str,
-
entry_point: Optional[Union[Callable, str]] = None,
-
mp_wrapper: RawInterfaceWrapper = DefaultMPWrapper,
-
register_step_based: bool = True, # TODO: Detect
-
add_mp_types: List[str] = KNOWN_MPS,
-
mp_config_override: Dict[str, Any] = {},
-
**kwargs
-
):
-
"""
-
Registers a Gymnasium environment, including Movement Primitives (MP) versions.
-
If you only want to register MP versions for an already registered environment, use fancy_gym.upgrade instead.
-
-
Args:
-
id (str): The unique identifier for the environment.
-
entry_point (Optional[Union[Callable, str]]): The entry point for creating the environment.
-
mp_wrapper (RawInterfaceWrapper): The MP wrapper for the environment.
-
register_step_based (bool): Whether to also register the raw srtep-based version of the environment (default True).
-
add_mp_types (List[str]): List of additional MP types to register.
-
mp_config_override (Dict[str, Any]): Dictionary for overriding MP configuration.
-
**kwargs: Additional keyword arguments which are passed to the environment constructor.
-
-
Notes:
-
- When `register_step_based` is True, the raw environment will also be registered to gymnasium otherwise only mp-versions will be registered.
-
- `entry_point` can be given as a string, allowing the same notation as gymnasium.
-
- If `id` already exists in the Gymnasium registry and `register_step_based` is True,
-
a warning message will be printed, suggesting to set `register_step_based=False` or use `fancy_gym.upgrade`.
-
-
Example:
-
To register a step-based environment with Movement Primitive versions (will use default mp_wrapper):
-
>>> register("MyEnv-v0", MyEnvClass"my_module:MyEnvClass")
-
-
The entry point can also be provided as a string:
-
>>> register("MyEnv-v0", "my_module:MyEnvClass")
-
-
"""
-
if register_step_based and id in gym_registry:
-
print(f'[Info] Gymnasium env with id "{id}" already exists. You should supply register_step_based=False or use fancy_gym.upgrade if you only want to register mp versions of an existing env.')
-
if register_step_based:
-
assert entry_point != None, 'You need to provide an entry-point, when registering step-based.'
-
if not callable(mp_wrapper): # mp_wrapper can be given as a String (same notation as for entry_point)
-
mod_name, attr_name = mp_wrapper.split(':')
-
mod = importlib.import_module(mod_name)
-
mp_wrapper = getattr(mod, attr_name)
-
if register_step_based:
-
gym_register(id=id, entry_point=entry_point, **kwargs)
-
upgrade(id, mp_wrapper, add_mp_types, mp_config_override)
-
-
-[docs]def upgrade(
-
id: str,
-
mp_wrapper: RawInterfaceWrapper = DefaultMPWrapper,
-
add_mp_types: List[str] = KNOWN_MPS,
-
base_id: Optional[str] = None,
-
mp_config_override: Dict[str, Any] = {},
-
):
-
"""
-
Upgrades an existing Gymnasium environment to include Movement Primitives (MP) versions.
-
We expect the raw step-based env to be already registered with gymnasium. Otherwise please use fancy_gym.register instead.
-
-
Args:
-
id (str): The unique identifier for the environment.
-
mp_wrapper (RawInterfaceWrapper): The MP wrapper for the environment (default is DefaultMPWrapper).
-
add_mp_types (List[str]): List of additional MP types to register (default is KNOWN_MPS).
-
base_id (Optional[str]): The unique identifier for the environment to upgrade. Will use id if non is provided. Can be defined to allow multiple registrations of different versions for the same step-based environment.
-
mp_config_override (Dict[str, Any]): Dictionary for overriding MP configuration.
-
-
Notes:
-
- The `id` parameter should match the ID of the existing Gymnasium environment you wish to upgrade. You can also pick a new one, but then `base_id` needs to be provided.
-
- The `mp_wrapper` parameter specifies the MP wrapper to use, allowing for customization.
-
- `add_mp_types` can be used to specify additional MP types to register alongside the base environment.
-
- The `base_id` parameter should match the ID of the existing Gymnasium environment you wish to upgrade.
-
- `mp_config_override` allows for customizing MP configuration if needed.
-
-
Example:
-
To upgrade an existing environment with MP versions:
-
>>> upgrade("MyEnv-v0", mp_wrapper=CustomMPWrapper)
-
-
To upgrade an existing environment with custom MP types and configuration:
-
>>> upgrade("MyEnv-v0", mp_wrapper=CustomMPWrapper, add_mp_types=["ProDMP", "DMP"], mp_config_override={"param": 42})
-
"""
-
if not base_id:
-
base_id = id
-
register_mps(id, base_id, mp_wrapper, add_mp_types, mp_config_override)
-
-
-def register_mps(id: str, base_id: str, mp_wrapper: RawInterfaceWrapper, add_mp_types: List[str] = KNOWN_MPS, mp_config_override: Dict[str, Any] = {}):
- for mp_type in add_mp_types:
- register_mp(id, base_id, mp_wrapper, mp_type, mp_config_override.get(mp_type, {}))
-
-
-def register_mp(id: str, base_id: str, mp_wrapper: RawInterfaceWrapper, mp_type: List[str], mp_config_override: Dict[str, Any] = {}):
- assert mp_type in KNOWN_MPS, 'Unknown mp_type'
- assert id not in ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS[mp_type], f'The environment {id} is already registered for {mp_type}.'
-
- parts = id.split('/')
- if len(parts) == 1:
- ns, name = 'gym', parts[0]
- elif len(parts) == 2:
- ns, name = parts[0], parts[1]
- else:
- raise ValueError('env id can not contain multiple "/".')
-
- parts = name.split('-')
- assert len(parts) >= 2 and parts[-1].startswith('v'), 'Malformed env id, must end in -v{int}.'
-
- fancy_id = f'{ns}_{mp_type}/{name}'
-
- gym_register(
- id=fancy_id,
- entry_point=bb_env_constructor,
- kwargs={
- 'underlying_id': base_id,
- 'mp_wrapper': mp_wrapper,
- 'mp_type': mp_type,
- '_mp_config_override_register': mp_config_override
- }
- )
-
- ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS[mp_type].append(fancy_id)
- ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all'].append(fancy_id)
- if ns not in MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS:
- MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS[ns] = {mp_type: [] for mp_type in _KNOWN_MPS_PLUS_ALL}
- MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS[ns][mp_type].append(fancy_id)
- MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS[ns]['all'].append(fancy_id)
-
-
-def nested_update(base: MutableMapping, update):
- """
- Updated method for nested Mappings
- Args:
- base: main Mapping to be updated
- update: updated values for base Mapping
-
- """
- if any([item.endswith('_type') for item in update]):
- base = update
- return base
- for k, v in update.items():
- base[k] = nested_update(base.get(k, {}), v) if isinstance(v, Mapping) else v
- return base
-
-
-def bb_env_constructor(underlying_id, mp_wrapper, mp_type, mp_config_override={}, _mp_config_override_register={}, **kwargs):
- raw_underlying_env = gym_make(underlying_id, **kwargs)
- underlying_env = mp_wrapper(raw_underlying_env)
-
- mp_config = getattr(underlying_env, 'mp_config') if hasattr(underlying_env, 'mp_config') else {}
- active_mp_config = copy.deepcopy(mp_config.get(mp_type, {}))
- global_inherit_defaults = mp_config.get('inherit_defaults', True)
- inherit_defaults = active_mp_config.pop('inherit_defaults', global_inherit_defaults)
-
- config = copy.deepcopy(_BB_DEFAULTS[mp_type]) if inherit_defaults else {}
- nested_update(config, active_mp_config)
- nested_update(config, _mp_config_override_register)
- nested_update(config, mp_config_override)
-
- wrappers = config.pop('wrappers')
-
- traj_gen_kwargs = config.pop('trajectory_generator_kwargs', {})
- black_box_kwargs = config.pop('black_box_kwargs', {})
- contr_kwargs = config.pop('controller_kwargs', {})
- phase_kwargs = config.pop('phase_generator_kwargs', {})
- basis_kwargs = config.pop('basis_generator_kwargs', {})
-
- return make_bb(underlying_env,
- wrappers=wrappers,
- black_box_kwargs=black_box_kwargs,
- traj_gen_kwargs=traj_gen_kwargs,
- controller_kwargs=contr_kwargs,
- phase_kwargs=phase_kwargs,
- basis_kwargs=basis_kwargs,
- **config)
-
- 
-
-
- 
- 
- 
