mani_skill.utils.wrappers#

Submodules#

Classes#

`ActionRepeatWrapper`	Wraps a `gymnasium.Env` to allow a modular transformation of the `step()` and `reset()` methods.
`CPUGymWrapper`	This wrapper wraps any maniskill env created via gym.make to ensure the outputs of
`CachedResetWrapper`	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.
`FlattenActionSpaceWrapper`	Flattens the action space. The original action space must be spaces.Dict
`FlattenObservationWrapper`	Flattens the observations into a single vector
`FlattenRGBDObservationWrapper`	Flattens the rgbd mode observations into a dictionary with two keys, "rgbd" and "state"
`FrameStack`	Observation wrapper that stacks the observations in a rolling manner.

Package Contents#

class mani_skill.utils.wrappers.ActionRepeatWrapper(env, repeat)[source]#

Bases: gymnasium.Wrapper

Wraps a gymnasium.Env to allow a modular transformation of the step() and 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 action_space, observation_space and metadata attributes, without changing the underlying environment’s attributes. Moreover, the behavior of the step() and reset() methods can be changed by these wrappers.

Some attributes (spec, render_mode, np_random) will point back to the wrapper’s environment (i.e. to the corresponding attributes of env).

Note

If you inherit from Wrapper, don’t forget to call super().__init__(env)

Parameters:

env (mani_skill.envs.sapien_env.BaseEnv)
repeat (int)

_update_dict_values(from_dict, to_dict, not_dones)[source]#

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.

Parameters:

from_dict (dict)
to_dict (dict)
not_dones (mani_skill.utils.structs.types.Array)

step(action)[source]#: Uses the step() of the env that can be overwritten to change the returned data.

property base_env: mani_skill.envs.sapien_env.BaseEnv#

Return type:: mani_skill.envs.sapien_env.BaseEnv

property num_envs#

repeat#

class mani_skill.utils.wrappers.CPUGymWrapper(env, ignore_terminations=False, record_metrics=False)[source]#

Bases: 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

Parameters:

env (gym.Env) – The environment to wrap.
ignore_terminations (bool) – If True, the environment will ignore termination signals and continue running until truncation. Default is False.
record_metrics (bool) – 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.

render()[source]#: Uses the render() of the env that can be overwritten to change the returned data.

reset(*, seed=None, options=None)[source]#: Uses the reset() of the env that can be overwritten to change the returned data.

step(action)[source]#: Uses the step() of the env that can be overwritten to change the returned data.

action_space#: Return the Env action_space unless overwritten then the wrapper action_space is used.

property base_env: mani_skill.envs.sapien_env.BaseEnv#

Return type:: mani_skill.envs.sapien_env.BaseEnv

ignore_terminations = False#

observation_space#: Return the Env observation_space unless overwritten then the wrapper observation_space is used.

record_metrics = False#

class mani_skill.utils.wrappers.CachedResetWrapper(env, reset_to_env_states=None, config=CachedResetsConfig())[source]#

Bases: 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.

Parameters:

env (gymnasium.Env) – The environment to wrap.
reset_to_env_states (Optional[dict]) – 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.
config (Union[CachedResetsConfig, dict]) – A dictionary or a CachedResetsConfig object that contains the configuration for the cached resets.

reset(*args, seed=None, options=None, **kwargs)[source]#

Uses the reset() of the env that can be overwritten to change the returned data.

Parameters:

seed (Optional[Union[int, list[int]]])
options (Optional[dict])

_cached_resets_env_states = None#

_num_cached_resets = 0#

property base_env: mani_skill.envs.sapien_env.BaseEnv#

Return type:: mani_skill.envs.sapien_env.BaseEnv

cached_resets_config#

num_envs#

class mani_skill.utils.wrappers.FlattenActionSpaceWrapper(env)[source]#

Bases: gymnasium.ActionWrapper

Flattens the action space. The original action space must be spaces.Dict

action(action)[source]#

Returns a modified action before step() is called.

Parameters:: action – The original step() actions
Returns:: The modified actions

_orig_single_action_space#

property base_env: mani_skill.envs.sapien_env.BaseEnv#

Return type:: mani_skill.envs.sapien_env.BaseEnv

single_action_space#

class mani_skill.utils.wrappers.FlattenObservationWrapper(env)[source]#

Bases: gymnasium.ObservationWrapper

Flattens the observations into a single vector

observation(observation)[source]#

Returns a modified observation.

Parameters:: observation – The env observation
Returns:: The modified observation

property base_env: mani_skill.envs.sapien_env.BaseEnv#

Return type:: mani_skill.envs.sapien_env.BaseEnv

class mani_skill.utils.wrappers.FlattenRGBDObservationWrapper(env, rgb=True, depth=True, state=True, sep_depth=True)[source]#

Bases: gymnasium.ObservationWrapper

Flattens the rgbd mode observations into a dictionary with two keys, “rgbd” and “state”

Parameters:

rgb (bool) – Whether to include rgb images in the observation
depth (bool) – Whether to include depth images in the observation
state (bool) – Whether to include state data in the observation
sep_depth (bool) – Whether to separate depth and rgb images in the observation. Default is True.

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.

observation(observation)[source]#

Returns a modified observation.

Parameters:: observation (dict) – The env observation
Returns:: The modified observation

base_env: mani_skill.envs.sapien_env.BaseEnv#

include_depth = True#

include_rgb = True#

include_state = True#

sep_depth = True#

class mani_skill.utils.wrappers.FrameStack(env, num_stack, lz4_compress=False)[source]#

Bases: 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 reset() is called, the frame buffer will be filled with the initial observation. I.e. the observation returned by reset() will consist of num_stack many identical frames.

Parameters:

env (gymnasium.Env)
num_stack (int)
lz4_compress (bool)

observation(observation)[source]#

Returns a modified observation.

Parameters:: observation – The env observation
Returns:: The modified observation

reset(seed=None, options=None)[source]#

Reset the environment with kwargs.

Parameters:

seed – The seed for the environment reset
options – The options for the environment reset

Returns:

The stacked observations

step(action)[source]#

Steps through the environment, appending the observation to the frame buffer.

Parameters:: action – The action to step through the environment with
Returns:: Stacked observations, reward, terminated, truncated, and information from the environment

property base_env: mani_skill.envs.sapien_env.BaseEnv#

Return type:: mani_skill.envs.sapien_env.BaseEnv

frames#

lz4_compress = False#

num_stack#

use_dict#