mani_skill.utils.wrappers ========================= .. py:module:: mani_skill.utils.wrappers Submodules ---------- .. toctree:: :maxdepth: 1 /api/mani_skill/utils/wrappers/action_repeat/index /api/mani_skill/utils/wrappers/cached_reset/index /api/mani_skill/utils/wrappers/flatten/index /api/mani_skill/utils/wrappers/frame_stack/index /api/mani_skill/utils/wrappers/gymnasium/index /api/mani_skill/utils/wrappers/record/index /api/mani_skill/utils/wrappers/visual_encoders/index Classes ------- .. autoapisummary:: mani_skill.utils.wrappers.ActionRepeatWrapper mani_skill.utils.wrappers.CPUGymWrapper mani_skill.utils.wrappers.CachedResetWrapper mani_skill.utils.wrappers.FlattenActionSpaceWrapper mani_skill.utils.wrappers.FlattenObservationWrapper mani_skill.utils.wrappers.FlattenRGBDObservationWrapper mani_skill.utils.wrappers.FrameStack Package Contents ---------------- .. py:class:: ActionRepeatWrapper(env, repeat) Bases: :py:obj:`gymnasium.Wrapper` Wraps a :class:`gymnasium.Env` to allow a modular transformation of the :meth:`step` and :meth:`reset` methods. This class is the base class of all wrappers to change the behavior of the underlying environment. Wrappers that inherit from this class can modify the :attr:`action_space`, :attr:`observation_space` and :attr:`metadata` attributes, without changing the underlying environment's attributes. Moreover, the behavior of the :meth:`step` and :meth:`reset` methods can be changed by these wrappers. Some attributes (:attr:`spec`, :attr:`render_mode`, :attr:`np_random`) will point back to the wrapper's environment (i.e. to the corresponding attributes of :attr:`env`). .. note:: If you inherit from :class:`Wrapper`, don't forget to call ``super().__init__(env)`` .. py:method:: _update_dict_values(from_dict, to_dict, not_dones) Recursively updates the values of a dictionary with the values from another dictionary but only for the envs that are not done. This allows us to update the observation and info dictionaries with new values only for the environments that are not done. If a sub-env becomes done, its future step data will be discarded since not_dones will be false for this sub-environment. Therefore the final observation/info will come from the true last step of the sub-env. .. py:method:: step(action) Uses the :meth:`step` of the :attr:`env` that can be overwritten to change the returned data. .. py:property:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:property:: num_envs .. py:attribute:: repeat .. py:class:: CPUGymWrapper(env, ignore_terminations = False, record_metrics = False) Bases: :py:obj:`gymnasium.Wrapper` This wrapper wraps any maniskill env created via gym.make to ensure the outputs of env.render, env.reset, env.step are all numpy arrays and are not batched. Essentially ensuring the environment conforms entirely to the standard gymnasium API https://gymnasium.farama.org/api/env/. The wrapper also optionally records standardized evaluation metrics like return and success. This wrapper should generally be applied after all other wrappers as most wrappers for ManiSkill assume data returned is a batched torch tensor :param env: The environment to wrap. :type env: gym.Env :param ignore_terminations: If True, the environment will ignore termination signals and continue running until truncation. Default is False. :type ignore_terminations: bool :param record_metrics: If True, the returned info objects will contain the metrics: return, length, success_once, success_at_end, fail_once, fail_at_end. success/fail metrics are recorded only when the environment has success/fail criteria. success/fail_at_end are recorded only when ignore_terminations is True. :type record_metrics: bool .. py:method:: render() Uses the :meth:`render` of the :attr:`env` that can be overwritten to change the returned data. .. py:method:: reset(*, seed=None, options=None) Uses the :meth:`reset` of the :attr:`env` that can be overwritten to change the returned data. .. py:method:: step(action) Uses the :meth:`step` of the :attr:`env` that can be overwritten to change the returned data. .. py:attribute:: action_space Return the :attr:`Env` :attr:`action_space` unless overwritten then the wrapper :attr:`action_space` is used. .. py:property:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:attribute:: ignore_terminations :value: False .. py:attribute:: observation_space Return the :attr:`Env` :attr:`observation_space` unless overwritten then the wrapper :attr:`observation_space` is used. .. py:attribute:: record_metrics :value: False .. py:class:: CachedResetWrapper(env, reset_to_env_states = None, config = CachedResetsConfig()) Bases: :py:obj:`gymnasium.Wrapper` Cached reset wrapper for ManiSkill3 environments. Caching resets allows you to skip slower parts of the reset function call and boost environment FPS as a result. :param env: The environment to wrap. :param reset_to_env_states: A dictionary with keys "env_states" and optionally "obs". "env_states" is a dictionary of environment states to reset to. "obs" contains the corresponding observations generated at those env states. If reset_to_env_states is not provided, the wrapper will sample reset states from the environment using the given seed. :param config: A dictionary or a `CachedResetsConfig` object that contains the configuration for the cached resets. .. py:method:: reset(*args, seed = None, options = None, **kwargs) Uses the :meth:`reset` of the :attr:`env` that can be overwritten to change the returned data. .. py:attribute:: _cached_resets_env_states :value: None .. py:attribute:: _num_cached_resets :value: 0 .. py:property:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:attribute:: cached_resets_config .. py:attribute:: num_envs .. py:class:: FlattenActionSpaceWrapper(env) Bases: :py:obj:`gymnasium.ActionWrapper` Flattens the action space. The original action space must be spaces.Dict .. py:method:: action(action) Returns a modified action before :meth:`step` is called. :param action: The original :meth:`step` actions :returns: The modified actions .. py:attribute:: _orig_single_action_space .. py:property:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:attribute:: single_action_space .. py:class:: FlattenObservationWrapper(env) Bases: :py:obj:`gymnasium.ObservationWrapper` Flattens the observations into a single vector .. py:method:: observation(observation) Returns a modified observation. :param observation: The :attr:`env` observation :returns: The modified observation .. py:property:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:class:: FlattenRGBDObservationWrapper(env, rgb=True, depth=True, state=True, sep_depth=True) Bases: :py:obj:`gymnasium.ObservationWrapper` Flattens the rgbd mode observations into a dictionary with two keys, "rgbd" and "state" :param rgb: Whether to include rgb images in the observation :type rgb: bool :param depth: Whether to include depth images in the observation :type depth: bool :param state: Whether to include state data in the observation :type state: bool :param sep_depth: Whether to separate depth and rgb images in the observation. Default is True. :type sep_depth: bool Note that the returned observations will have a "rgb" or "depth" key depending on the rgb/depth bool flags, and will always have a "state" key. If sep_depth is False, rgb and depth will be merged into a single "rgbd" key. .. py:method:: observation(observation) Returns a modified observation. :param observation: The :attr:`env` observation :returns: The modified observation .. py:attribute:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:attribute:: include_depth :value: True .. py:attribute:: include_rgb :value: True .. py:attribute:: include_state :value: True .. py:attribute:: sep_depth :value: True .. py:class:: FrameStack(env, num_stack, lz4_compress = False) Bases: :py:obj:`gymnasium.ObservationWrapper` Observation wrapper that stacks the observations in a rolling manner. For example, if the number of stacks is 4, then the returned observation contains the most recent 4 observations. For environment 'PickCube-v1', the original observation is an array with shape [42], so if we stack 4 observations, the processed observation has shape [4, 42]. This wrapper also supports dict observations, and will stack the leafs of the dictionary accordingly. .. note:: - After :meth:`reset` is called, the frame buffer will be filled with the initial observation. I.e. the observation returned by :meth:`reset` will consist of `num_stack` many identical frames. .. py:method:: observation(observation) Returns a modified observation. :param observation: The :attr:`env` observation :returns: The modified observation .. py:method:: reset(seed=None, options=None) Reset the environment with kwargs. :param seed: The seed for the environment reset :param options: The options for the environment reset :returns: The stacked observations .. py:method:: step(action) Steps through the environment, appending the observation to the frame buffer. :param action: The action to step through the environment with :returns: Stacked observations, reward, terminated, truncated, and information from the environment .. py:property:: base_env :type: mani_skill.envs.sapien_env.BaseEnv .. py:attribute:: frames .. py:attribute:: lz4_compress :value: False .. py:attribute:: num_stack .. py:attribute:: use_dict